Ringkasan
Halaman ini menjelaskan desain dan langkah-langkah yang diperlukan untuk mengonversi operasi gabungan di TensorFlow menjadi operasi gabungan di TensorFlow Lite. Infrastruktur ini bertujuan umum dan mendukung konversi operasi gabungan apa pun di TensorFlow menjadi operasi gabungan yang sesuai di TensorFlow Lite.
Contoh penggunaan infrastruktur ini adalah penggabungan operasi TensorFlow RNN ke TensorFlow Lite, seperti yang dirinci di sini .
Apa itu operasi gabungan
Operasi TensorFlow dapat berupa operasi primitif misalnya tf.add atau dapat terdiri dari operasi primitif lainnya misalnya tf.einsum . Operasi primitif muncul sebagai satu node di grafik TensorFlow, sedangkan operasi gabungan adalah kumpulan node di grafik TensorFlow. Mengeksekusi operasi gabungan sama dengan mengeksekusi setiap operasi primitif penyusunnya.
Operasi gabungan berhubungan dengan operasi tunggal yang menggabungkan semua komputasi yang dilakukan oleh setiap operasi primitif dalam operasi gabungan terkait.
Manfaat operasi gabungan
Operasi gabungan ada untuk memaksimalkan kinerja implementasi kernel yang mendasarinya, dengan mengoptimalkan keseluruhan komputasi dan mengurangi jejak memori. Hal ini sangat berharga, terutama untuk beban kerja inferensi latensi rendah dan platform seluler dengan sumber daya terbatas.
Operasi gabungan juga menyediakan antarmuka tingkat yang lebih tinggi untuk mendefinisikan transformasi kompleks seperti kuantisasi, yang jika tidak dilakukan akan menjadi tidak mungkin atau sangat sulit dilakukan pada tingkat yang lebih terperinci.
TensorFlow Lite memiliki banyak contoh operasi gabungan karena alasan yang disebutkan di atas. Operasi gabungan ini biasanya berhubungan dengan operasi gabungan dalam program sumber TensorFlow. Contoh operasi komposit di TensorFlow yang diimplementasikan sebagai operasi gabungan tunggal di TensorFlow Lite mencakup berbagai operasi RNN seperti LSTM rangkaian Searah dan Dua Arah, konvolusi (konv2d, penambahan bias, relu), terhubung sepenuhnya (matmul, penambahan bias, relu) dan banyak lagi . Di TensorFlow Lite, kuantisasi LSTM saat ini hanya diterapkan dalam operasi LSTM yang digabungkan.
Tantangan dengan operasi yang menyatu
Mengonversi operasi gabungan dari TensorFlow menjadi operasi gabungan di TensorFlow Lite adalah masalah yang sulit. Hal ini dikarenakan:
Operasi gabungan direpresentasikan dalam grafik TensorFlow sebagai serangkaian operasi primitif tanpa batas yang jelas. Akan sangat sulit untuk mengidentifikasi (misalnya melalui pencocokan pola) sub-grafik yang terkait dengan operasi gabungan 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 implementasi ini terdiri dari operasi primitif yang berbeda, tetapi semuanya masih dapat dikonversi ke operasi LSTM gabungan yang sama di TensorFlow Lite.
Oleh karena itu, 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 model dapat dipetakan ke satu operasi di TFLite. Hal ini dapat membantu kinerja saat menulis implementasi yang dioptimalkan untuk operasi tertentu. Untuk dapat membuat operasi gabungan di TFLite, identifikasi bagian grafik yang mewakili operasi gabungan dan gabungkan dalam tf.function
dengan atribut "experimental_implements" ke tf.function
, yang memiliki nilai atribut tfl_fusable_op
dengan nilai true
. Jika operasi khusus mengambil atribut, teruskan atribut tersebut sebagai bagian dari "implementasi_eksperimental" 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 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 operasi harus serupa dengan nama yang ditentukan dalam atribut name
di tanda tangan implementasi.
Contoh op pada contoh tersebut 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 komposit ke operasi leburan (Lanjutan)
Arsitektur keseluruhan untuk mengonversi operasi gabungan TensorFlow ke operasi gabungan TensorFlow Lite adalah sebagai berikut:
Bungkus operasi komposit dalam tf.function
Dalam kode sumber model TensorFlow, identifikasi dan abstraksi operasi komposit menjadi tf.function
dengan anotasi fungsi eksperimental_implements . Lihat contoh penyematan pencarian . Fungsi tersebut 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 antarmuka ini dengan implementasi gabungan.
Di pass fungsi persiapan-komposit, masukkan kode konversi Anda.
Dalam penggunaan yang lebih maju, dimungkinkan untuk menerapkan transformasi kompleks dari operan operasi gabungan untuk mendapatkan operan dari operasi gabungan. 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 operasi gabungan di TensorFlow Lite.
Menulis operasi 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 terdefinisi dengan baik untuk operasi gabungan dalam grafik TensorFlow yang mendasarinya.
- Tentukan secara eksplisit antarmuka yang diimplementasikan oleh operasi ini. Argumen dari
tf.function
sesuai dengan argumen antarmuka ini.
Sebagai contoh, mari kita pertimbangkan operasi gabungan yang ditentukan untuk mengimplementasikan pencarian penyematan. 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, infrastruktur umum dapat dibangun 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 impor model TensorFlow sebagai grafik dengan semua variabel diganti dengan nilai konstanta yang sesuai. Ini tidak berlaku untuk operasi fusi karena grafik tersebut memiliki semua fungsi yang sejajar sehingga variabel dapat diubah menjadi konstanta.
Untuk memanfaatkan tf.function
dengan fitur experimental_implements
selama proses konversi, fungsi-fungsi tersebut perlu dipertahankan hingga proses konversi nanti.
Oleh 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-fitur baru yang ditambahkan adalah:
- Mengimpor model tersimpan TensorFlow ke MLIR
- sekering operasi komposit
- analisis mutabilitas variabel
- bekukan semua variabel read-only
Hal ini memungkinkan kita untuk melakukan fusi operasi menggunakan fungsi yang mewakili operasi komposit sebelum fungsi inlining dan pembekuan variabel.
Menerapkan fusi operasi
Mari kita lihat operasi 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, panggil utilitas fusi operasi yang sesuai.
- Utilitas fusi operasi beroperasi pada operan dan atribut fungsi (yang berfungsi sebagai antarmuka untuk konversi) dan menggantikan isi fungsi dengan isi 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 gabungan. Karena semua penghitungan ini dapat dilipat secara konstan, penghitungan tersebut tidak akan ada di flatbuffer yang diekspor di mana hanya operasi gabungan yang ada.
Berikut 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 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());
}