Ringkasan
Halaman ini menjelaskan desain dan langkah-langkah yang diperlukan untuk mengonversi operasi gabungan di TensorFlow menjadi operasi gabungan di TensorFlow Lite. Infrastruktur ini bersifat umum dan mendukung konversi operasi gabungan apa pun di TensorFlow ke operasi fusi yang sesuai di TensorFlow Lite.
Contoh penggunaan infrastruktur ini adalah fusi operasi TensorFlow RNN ke TensorFlow Lite, seperti yang dijelaskan di sini.
Apa itu operasi gabungan
Operasi TensorFlow dapat berupa operasi primitif, misalnya tf.add atau dapat dikomposisi dari operasi primitif lainnya, misalnya tf.einsum. Operasi primitif ditampilkan sebagai node tunggal dalam grafik TensorFlow, sedangkan operasi gabungan adalah kumpulan node dalam grafik TensorFlow. Mengeksekusi operasi gabungan sama dengan mengeksekusi setiap operasi primitif konstituennya.
Operasi gabungan berkaitan dengan satu operasi yang memasukkan semua komputasi yang dilakukan oleh setiap operasi primitif dalam operasi gabungan yang sesuai.
Manfaat operasi gabungan
Operasi gabungan diterapkan untuk memaksimalkan performa implementasi kernel yang mendasarinya, dengan mengoptimalkan komputasi keseluruhan dan mengurangi jejak memori. Hal ini sangat berharga, terutama untuk beban kerja inferensi latensi rendah dan platform seluler dengan resource yang terbatas.
Operasi gabungan juga menyediakan antarmuka tingkat yang lebih tinggi untuk menentukan transformasi kompleks seperti kuantisasi, yang tidak mungkin atau sangat sulit dilakukan pada tingkat yang lebih terperinci.
TensorFlow Lite memiliki banyak instance operasi gabungan karena alasan yang disebutkan di atas. Operasi gabungan ini biasanya berkaitan dengan operasi gabungan dalam program TensorFlow sumber. Contoh operasi gabungan di TensorFlow yang diimplementasikan sebagai operasi gabungan tunggal di TensorFlow Lite mencakup berbagai operasi RNN seperti Searah Unidirectional dan Bidirectional, LSTM, konvolusi (konv2d, penambahan bias, relu), terhubung sepenuhnya (matmul, penambahan bias, relu) dan banyak lagi. Di TensorFlow Lite, kuantisasi LSTM saat ini hanya diimplementasikan dalam operasi LSTM fusi.
Tantangan dalam operasi gabungan
Mengonversi operasi gabungan dari TensorFlow ke operasi gabungan di TensorFlow Lite adalah masalah yang sulit. Hal ini dikarenakan:
Operasi gabungan direpresentasikan dalam grafik TensorFlow sebagai sekumpulan operasi primitif tanpa batas yang ditentukan dengan baik. Akan sangat sulit untuk mengidentifikasi (misalnya melalui pencocokan pola) sub-grafik yang terkait dengan operasi komposit tersebut.
Mungkin ada lebih dari satu implementasi TensorFlow yang menargetkan operasi TensorFlow Lite gabungan. Misalnya, ada banyak implementasi LSTM di TensorFlow (Keras, Babelfish/lingvo, dll.) dan masing-masing terdiri dari operasi primitif yang berbeda, tetapi semuanya masih dapat dikonversi ke operasi LSTM fusi yang sama di TensorFlow Lite.
Dengan demikian, konversi operasi gabungan terbukti cukup sulit.
Mengonversi dari operasi gabungan ke operasi khusus TFLite (disarankan)
Menggabungkan operasi gabungan dalam tf.function
Dalam banyak kasus, beberapa bagian model dapat dipetakan ke satu operasi di TFLite. Hal ini dapat membantu meningkatkan performa saat menulis implementasi yang dioptimalkan
untuk operasi tertentu. Agar dapat membuat operasi gabungan di TFLite,
identifikasi bagian grafik yang merepresentasikan operasi gabungan, lalu gabungkan dalam
tf.function dengan
atribut "experimental_implementations" ke tf.function, yang memiliki nilai
atribut tfl_fusable_op dengan nilai true. Jika operasi kustom mengambil atribut, teruskan atribut sebagai bagian dari "experimental_implementations" yang sama.
Contoh:
def get_implements_signature():
implements_signature = [
# 'name' will be used as a name for the operation.
'name: "my_custom_fused_op"',
# attr "tfl_fusable_op" is required to be set with true value.
'attr {key: "tfl_fusable_op" value { b: true } }',
# Example attribute "example_option" that the op accepts.
'attr {key: "example_option" value { i: %d } }' % 10
]
return ' '.join(implements_signature)
@tf.function(experimental_implements=get_implements_signature())
def my_custom_fused_op(input_1, input_2):
# An empty function that represents pre/post processing example that
# is not represented as part of the Tensorflow graph.
output_1 = tf.constant(0.0, dtype=tf.float32, name='first_output')
output_2 = tf.constant(0.0, dtype=tf.float32, name='second_output')
return output_1, output_2
class TestModel(tf.Module):
def __init__(self):
super(TestModel, self).__init__()
self.conv_1 = tf.keras.layers.Conv2D(filters=1, kernel_size=(3, 3))
self.conv_2 = tf.keras.layers.Conv2D(filters=1, kernel_size=(3, 3))
@tf.function(input_signature=[
tf.TensorSpec(shape=[1, 28, 28, 3], dtype=tf.float32),
tf.TensorSpec(shape=[1, 28, 28, 3], dtype=tf.float32),
])
def simple_eval(self, input_a, input_b):
return my_custom_fused_op(self.conv_1(input_a), self.conv_2(input_b))
Perhatikan bahwa Anda tidak perlu menetapkan allow_custom_ops pada pengonversi karena
atribut tfl_fusable_op sudah menyiratkan hal ini.
Implementasikan pengoperasian kustom dan daftar ke TFLite Interpreter
Implementasikan operasi fusi Anda sebagai operasi TFLite Custom - lihat instructions.
Perhatikan bahwa nama yang akan digunakan untuk mendaftarkan operasi harus mirip dengan nama
yang ditentukan dalam atribut name di tanda tangan implementasi.
Contoh untuk pengoperasian dalam contoh ini adalah
TfLiteRegistration reg = {};
// This name must match the name specified in the implements signature.
static constexpr char kOpName[] = "my_custom_fused_op";
reg.custom_name = kOpName;
reg.prepare = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
// Add your code.
return kTfLiteOk;
};
reg.invoke = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
// Add your code.
return kTfLiteOk;
};
reg.builtin_code = kTfLiteCustom;
resolver->AddCustom(kOpName, ®);
Mengonversi dari operasi gabungan ke operasi gabungan (Lanjutan)
Keseluruhan arsitektur untuk mengonversi operasi komposit TensorFlow ke operasi gabungan TensorFlow Lite adalah sebagai berikut:
Menggabungkan operasi gabungan dalam tf.function
Dalam kode sumber model TensorFlow, identifikasi dan pisahkan operasi gabungan menjadi tf.function dengan anotasi fungsi experimental_implements. Lihat contoh pencarian penyematan. Fungsi ini menentukan antarmuka dan argumennya harus digunakan untuk mengimplementasikan logika konversi.
Menulis kode konversi
Kode konversi ditulis per antarmuka fungsi dengan
anotasi implements. Lihat contoh penggabungan untuk menyematkan
pencarian. Secara konseptual, kode konversi menggantikan penerapan gabungan antarmuka ini dengan penerapan gabungan.
Di pass prepare-composite-function, tambahkan kode konversi Anda.
Pada penggunaan yang lebih canggih, Anda dapat mengimplementasikan transformasi kompleks dari operand operasi gabungan untuk memperoleh operand operasi gabungan. Lihat kode konversi Keras LSTM. sebagai contoh.
Konversikan ke TensorFlow Lite
Gunakan API TFLiteConverter.from_saved_model untuk melakukan konversi ke TensorFlow Lite.
Di balik layar
Kini kami menjelaskan detail keseluruhan desain dalam melakukan konversi ke operasi gabungan di TensorFlow Lite.
Menulis operasi di TensorFlow
Penggunaan tf.function dengan atribut fungsi experimental_implements memungkinkan pengguna untuk menyusun operasi baru secara eksplisit menggunakan operasi primitif TensorFlow dan menentukan antarmuka yang diimplementasikan oleh operasi komposit yang dihasilkan. Hal ini sangat berguna karena memberikan:
- Batas yang jelas untuk operasi gabungan dalam grafik TensorFlow yang mendasarinya.
- Tentukan secara eksplisit antarmuka yang diimplementasikan operasi ini. Argumen tf.function sesuai dengan argumen antarmuka ini.
Sebagai contoh, mari kita pertimbangkan operasi gabungan yang ditentukan untuk mengimplementasikan pencarian embedding. Objek ini dipetakan ke operasi gabungan di TensorFlow Lite.
@tf.function(
experimental_implements="embedding_lookup")
def EmbFprop(embs, ids_vec):
"""Embedding forward prop.
Effectively, it computes:
num = size of ids_vec
rets = zeros([num, embedding dim])
for i in range(num):
rets[i, :] = embs[ids_vec[i], :]
return rets
Args:
embs: The embedding matrix.
ids_vec: A vector of int32 embedding ids.
Returns:
The result of embedding lookups. A matrix of shape
[num ids in ids_vec, embedding dims].
"""
num = tf.shape(ids_vec)[0]
rets = inplace_ops.empty([num] + emb_shape_suf, py_utils.FPropDtype(p))
def EmbFpropLoop(i, embs, ids_vec, rets):
# row_id = ids_vec[i]
row_id = tf.gather(ids_vec, i)
# row = embs[row_id]
row = tf.reshape(tf.gather(embs, row_id), [1] + emb_shape_suf)
# rets[i] = row
rets = inplace_ops.alias_inplace_update(rets, [i], row)
return embs, ids_vec, rets
_, _, rets = functional_ops.For(
start=0,
limit=num,
delta=1,
inputs=[embs, ids_vec, rets],
body=EmbFpropLoop,
rewrite_with_while=compiled)
if len(weight_shape) > 2:
rets = tf.reshape(rets, [num, symbolic.ToStatic(p.embedding_dim)])
return rets
Dengan membuat model menggunakan operasi komposit melalui tf.function seperti yang diilustrasikan di atas, membangun infrastruktur umum untuk mengidentifikasi dan mengonversi operasi tersebut menjadi operasi TensorFlow Lite gabungan dapat dilakukan.
Memperluas konverter TensorFlow Lite
Konverter TensorFlow Lite yang dirilis awal tahun ini hanya mendukung pengimporan model TensorFlow sebagai grafik dengan semua variabel yang diganti dengan nilai konstantanya yang sesuai. Ini tidak berfungsi untuk penggabungan operasi karena grafik tersebut memiliki semua fungsi yang inline sehingga variabel dapat diubah menjadi konstanta.
Untuk memanfaatkan tf.function dengan fitur experimental_implements selama proses konversi, fungsi tersebut harus dipertahankan sampai nanti dalam proses konversi.
Oleh karena itu, kami menerapkan alur kerja baru untuk mengimpor dan mengonversi model TensorFlow di konverter untuk mendukung kasus penggunaan fusi operasi gabungan. Secara khusus, fitur baru yang ditambahkan adalah:
- Mengimpor model tersimpan ke MLIR TensorFlow
- operasi gabungan sekring
- analisis mutabilitas variabel
- membekukan semua variabel hanya baca
Hal ini memungkinkan kami melakukan fusi operasi menggunakan fungsi yang mewakili operasi gabungan sebelum fungsi inline dan pembekuan variabel.
Mengimplementasikan fusion operasi
Mari kita lihat kartu fusi operasi secara lebih mendetail. Kartu ini akan melakukan hal berikut:
- Melakukan loop semua fungsi dalam modul MLIR.
- Jika sebuah fungsi memiliki atribut tf._implementations, berdasarkan nilai atribut, akan memanggil utilitas fusi operasi yang sesuai.
- Utilitas fusi operasi beroperasi pada operand dan atribut fungsi (yang berfungsi sebagai antarmuka untuk konversi) dan menggantikan isi fungsi dengan isi fungsi setara yang berisi operasi gabungan.
- Dalam banyak kasus, isi yang diganti akan berisi operasi selain operasi gabungan. Ini berkaitan dengan beberapa transformasi statis pada Operand fungsi untuk mendapatkan operand dari operasi gabungan. Karena semua komputasi ini dapat dilipat secara konstan, komputasi tersebut tidak akan ada dalam flatbuffer yang diekspor di mana hanya operasi gabungan yang akan ada.
Berikut adalah cuplikan kode dari kartu yang menunjukkan alur kerja utama:
void PrepareCompositeFunctionsPass::ConvertTFImplements(FuncOp func,
StringAttr attr) {
if (attr.getValue() == "embedding_lookup") {
func.eraseBody();
func.addEntryBlock();
// Convert the composite embedding_lookup function body to a
// TFLite fused embedding_lookup op.
ConvertEmbeddedLookupFunc convert_embedded_lookup(func);
if (failed(convert_embedded_lookup.VerifySignature())) {
return signalPassFailure();
}
convert_embedded_lookup.RewriteFunc();
} else if (attr.getValue() == mlir::TFL::kKerasLstm) {
func.eraseBody();
func.addEntryBlock();
OpBuilder builder(func.getBody());
if (failed(ConvertKerasLSTMLayer(func, &builder))) {
return signalPassFailure();
}
} else if (.....) /* Other fusions can plug in here */
}
Berikut adalah cuplikan kode yang menunjukkan pemetaan operasi gabungan ini ke operasi gabungan di TensorFlow Lite yang memanfaatkan fungsi sebagai antarmuka konversi.
void RewriteFunc() {
Value lookup = func_.getArgument(1);
Value value = func_.getArgument(0);
auto output_type = func_.getType().getResult(0);
OpBuilder builder(func_.getBody());
auto op = builder.create<mlir::TFL::EmbeddingLookupOp>(
func_.getLoc(), output_type, lookup, value);
builder.create<mlir::ReturnOp>(func_.getLoc(), op.getResult());
}