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 gabungan yang sesuai di TensorFlow Lite.
Contoh penggunaan infrastruktur ini adalah penggabungan operasi TensorFlow RNN ke TensorFlow Lite, seperti yang dijelaskan di sini .
Apa itu operasi fusi?
Operasi TensorFlow dapat berupa operasi primitif misalnya tf.add atau dapat dibuat dari operasi primitif lainnya misalnya tf.einsum . Operasi primitif muncul sebagai node tunggal dalam grafik TensorFlow sedangkan operasi gabungan adalah kumpulan node dalam grafik TensorFlow. Mengeksekusi operasi komposit setara dengan mengeksekusi setiap operasi primitif penyusunnya.
Operasi fusi sesuai dengan operasi tunggal yang memasukkan semua perhitungan yang dilakukan oleh setiap operasi primitif dalam operasi komposit yang sesuai.
Manfaat operasi fusi
Operasi fusi ada untuk memaksimalkan kinerja implementasi kernel yang mendasarinya, dengan mengoptimalkan perhitungan keseluruhan dan mengurangi jejak memori. Ini sangat berharga, terutama untuk beban kerja inferensi latensi rendah dan platform seluler dengan sumber daya terbatas.
Operasi fusi juga menyediakan antarmuka tingkat yang lebih tinggi untuk mendefinisikan transformasi kompleks seperti kuantisasi, yang jika tidak, tidak mungkin atau sangat sulit dilakukan pada tingkat yang lebih terperinci.
TensorFlow Lite memiliki banyak contoh operasi gabungan untuk alasan yang disebutkan di atas. Operasi gabungan ini biasanya sesuai dengan operasi gabungan dalam program TensorFlow sumber. Contoh operasi gabungan di TensorFlow yang diimplementasikan sebagai operasi fusi tunggal di TensorFlow Lite mencakup berbagai operasi RNN seperti LSTM urutan searah dan dua arah, konvolusi (conv2d, bias add, relu), sepenuhnya terhubung (matmul, bias add, relu) dan banyak lagi . Di TensorFlow Lite, kuantisasi LSTM saat ini hanya diterapkan dalam operasi LSTM yang digabungkan.
Tantangan dengan operasi fusi
Mengonversi operasi komposit dari TensorFlow ke operasi gabungan di TensorFlow Lite adalah masalah yang sulit. Hal ini karena:
Operasi gabungan direpresentasikan dalam grafik TensorFlow sebagai kumpulan operasi primitif tanpa batas yang jelas. Ini bisa sangat menantang untuk mengidentifikasi (misalnya melalui pencocokan pola) sub-grafik yang sesuai dengan operasi komposit tersebut.
Mungkin ada lebih dari satu implementasi TensorFlow yang menargetkan operasi TensorFlow Lite yang digabungkan. 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 fusi LSTM yang sama di TensorFlow Lite.
Dengan demikian, konversi operasi fusi terbukti cukup menantang.
Mengonversi dari operasi komposit ke operasi kustom TFLite (disarankan)
Bungkus operasi komposit dalam tf.function
Dalam banyak kasus, beberapa bagian dari model dapat dipetakan ke satu operasi di TFLite. Ini dapat membantu kinerja saat menulis implementasi yang dioptimalkan untuk operasi tertentu. Untuk dapat membuat operasi fusi di TFLite, identifikasi bagian grafik yang mewakili operasi fusi dan bungkus dalam tf.function dengan atribut "experimental_implements" ke tf.function , yang memiliki nilai atribut tfl_fusable_op dengan nilai true . Jika operasi kustom mengambil atribut kemudian meneruskannya sebagai bagian dari "experimental_implements" 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 menyetel allow_custom_ops pada konverter karena atribut tfl_fusable_op sudah menyiratkan hal ini.
Terapkan operasi khusus dan daftar dengan TFLite Interpreter
Terapkan operasi gabungan Anda sebagai operasi Kustom TFLite - lihat instruksi .
Perhatikan bahwa, nama untuk mendaftarkan op harus mirip dengan nama yang ditentukan dalam atribut name di tanda tangan implements.
Contoh untuk op dalam contoh 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 coder.
return kTfLiteOk;
};
reg.builtin_code = kTfLiteCustom;
resolver->AddCustom(kOpName, ®);
Mengonversi dari operasi komposit ke operasi gabungan (Lanjutan)
Arsitektur keseluruhan untuk mengonversi operasi komposit TensorFlow ke operasi gabungan TensorFlow Lite di bawah ini:
Bungkus operasi komposit dalam tf.function
Dalam kode sumber model TensorFlow, identifikasi dan abstraksi operasi gabungan menjadi tf.function dengan anotasi fungsi eksperimental_implements . Lihat contoh embedding lookup . Fungsi mendefinisikan antarmuka dan argumennya harus digunakan untuk mengimplementasikan logika konversi.
Tulis kode konversi
Kode konversi ditulis per antarmuka fungsi dengan anotasi implements . Lihat contoh fusi untuk menyematkan pencarian . Secara konseptual, kode konversi menggantikan implementasi gabungan dari antarmuka ini dengan yang menyatu.
Di pass prep-composite-functions, plugin di kode konversi Anda .
Dalam penggunaan yang lebih maju, dimungkinkan untuk mengimplementasikan transformasi kompleks dari operan operasi komposit untuk mendapatkan operan dari operasi fusi. Lihat Keras LSTM . kode konversi sebagai contoh.
Konversikan ke TensorFlow Lite
Gunakan TFLiteConverter.from_saved_model API untuk mengonversi ke TensorFlow Lite.
Dibawah tenda
Kami sekarang menjelaskan detail tingkat tinggi dari keseluruhan desain dalam mengonversi ke operasi gabungan di TensorFlow Lite.
Operasi penulisan di TensorFlow
Penggunaan tf.function dengan atribut fungsi eksperimental_implements memungkinkan pengguna membuat operasi baru secara eksplisit menggunakan operasi primitif TensorFlow dan menentukan antarmuka yang diimplementasikan oleh operasi gabungan yang dihasilkan. Ini sangat berguna karena menyediakan:
- Batas yang jelas untuk operasi gabungan dalam grafik TensorFlow yang mendasarinya.
- Tentukan secara eksplisit antarmuka yang diimplementasikan oleh operasi ini. Argumen
tf.functionsesuai dengan argumen antarmuka ini.
Sebagai contoh, mari pertimbangkan operasi gabungan yang ditentukan untuk mengimplementasikan pencarian penyematan. Ini memetakan 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, menjadi mungkin untuk membangun infrastruktur umum untuk mengidentifikasi dan mengonversi operasi tersebut menjadi operasi TensorFlow Lite yang digabungkan.
Memperluas konverter TensorFlow Lite
Konverter TensorFlow Lite yang dirilis awal tahun ini hanya mendukung pengimporan model TensorFlow sebagai grafik dengan semua variabel diganti dengan nilai konstanta yang sesuai. Ini tidak bekerja untuk operasi fusi karena grafik tersebut memiliki semua fungsi yang digariskan sehingga variabel dapat diubah menjadi konstanta.
Untuk memanfaatkan tf.function dengan fitur experimental_implements selama proses konversi, fungsi tersebut perlu dipertahankan hingga nanti dalam proses konversi.
Karena itu, kami menerapkan alur kerja baru untuk mengimpor dan mengonversi model TensorFlow di konverter untuk mendukung kasus penggunaan fusi operasi komposit. Secara khusus, fitur baru yang ditambahkan adalah:
- Mengimpor model tersimpan TensorFlow ke MLIR
- operasi komposit sekering
- analisis mutabilitas variabel
- bekukan semua variabel hanya-baca
Hal ini memungkinkan kita untuk melakukan penggabungan operasi menggunakan fungsi yang mewakili operasi gabungan sebelum inlining fungsi dan pembekuan variabel.
Menerapkan operasi fusi
Mari kita lihat operan fusion pass lebih detail. Pass ini melakukan hal berikut:
- Ulangi semua fungsi dalam modul MLIR.
- Jika suatu fungsi memiliki atribut tf._implements, berdasarkan nilai atribut, memanggil utilitas fusi operasi yang sesuai.
- Utilitas fusi operasi beroperasi pada operan dan atribut fungsi (yang berfungsi sebagai antarmuka untuk konversi) dan menggantikan badan fungsi dengan badan fungsi setara yang berisi operasi fusi.
- Dalam banyak kasus, bodi yang diganti akan berisi operasi selain operasi fusi. Ini sesuai dengan beberapa transformasi statis pada operan fungsi untuk mendapatkan operan dari operasi fusi. Karena semua perhitungan ini dapat terus-menerus dilipat, mereka tidak akan ada di flatbuffer yang diekspor di mana hanya operasi gabungan yang akan ada.
Berikut adalah cuplikan kode dari pass 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());
}