Ringkasan
Halaman ini menjelaskan desain dan langkah-langkah yang diperlukan untuk mengonversi operasi gabungan di TensorFlow ke operasi gabungan di LiteRT. Infrastruktur ini tujuan umum dan mendukung konversi operasi gabungan apa pun di TensorFlow ke operasi fusi yang sesuai di LiteRT.
Contoh penggunaan infrastruktur ini adalah fusi operasi RNN TensorFlow untuk LiteRT, seperti yang dijelaskan di sini.
Apa itu operasi fusi
Operasi TensorFlow bisa berupa operasi primitif, misalnya tf.add atau dapat berupa yang terdiri dari operasi primitif lainnya, misalnya tf.einsum. Primitif muncul sebagai node tunggal pada grafik TensorFlow saat adalah kumpulan node dalam grafik TensorFlow. Menjalankan operasi gabungan setara dengan mengeksekusi setiap primitif konstituennya operasional bisnis.
Sebuah operasi gabungan sesuai dengan satu operasi yang komputasi yang dilakukan oleh setiap operasi primitif dalam setiap operasi gabungan.
Manfaat operasi gabungan
Operasi gabungan tersedia untuk memaksimalkan performa {i>kernel<i} yang mendasarinya dengan mengoptimalkan komputasi keseluruhan dan mengurangi memori jejak pelindung. Hal ini sangat bermanfaat, terutama untuk workload inferensi berlatensi rendah dan platform seluler dengan sumber daya terbatas.
Operasi gabungan juga menyediakan antarmuka dengan tingkat yang lebih tinggi untuk mendefinisikan transformasi seperti kuantisasi, yang seharusnya tidak layak dilakukan atau sangat sulit untuk dilakukan pada tingkat yang lebih terperinci.
LiteRT memiliki banyak instance operasi gabungan karena alasan yang diartikulasikan di atas. Operasi gabungan ini biasanya berkaitan dengan di program TensorFlow sumber. Contoh operasi gabungan dalam TensorFlow yang diimplementasikan sebagai operasi gabungan tunggal di LiteRT mencakup berbagai operasi RNN seperti urutan Searah dan Dua Arah LSTM, konvolusi (konv2d, bias tambah, relu), terhubung sepenuhnya (matmul, bias tambah, relu) dan banyak lagi. Di LiteRT, kuantisasi LSTM saat ini hanya diimplementasikan dalam operasi LSTM gabungan.
Tantangan dengan operasi gabungan
Mengubah operasi gabungan dari TensorFlow ke operasi gabungan di LiteRT adalah masalah yang sulit. Hal ini dikarenakan:
Operasi gabungan direpresentasikan dalam grafik TensorFlow sebagai serangkaian operasi primitif tanpa batasan yang jelas. Hal ini bisa 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 fused Operasi LiteRT. 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 menjadi operasi LSTM gabungan yang sama di LiteRT.
Dengan demikian, konversi operasi gabungan telah terbukti cukup menantang.
Mengonversi dari operasi gabungan ke operasi kustom TFLite (direkomendasikan)
Menggabungkan operasi gabungan dalam tf.function
Dalam banyak kasus, beberapa bagian model bisa dipetakan ke satu operasi di
TFLite. Hal ini dapat membantu meningkatkan performa saat menulis penerapan yang dioptimalkan
untuk operasi tertentu. Untuk dapat membuat operasi
fusi di TFLite,
mengidentifikasi bagian grafik yang merepresentasikan operasi menyatu dan menggabungkannya
tf.function dengan
"eksperimen_implementasi" ke tf.function, yang memiliki atribut
nilai tfl_fusable_op dengan nilai true. Jika operasi khusus mengambil
lalu meneruskannya 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 konverter sebagai
Atribut tfl_fusable_op sudah menyiratkan ini.
Terapkan operasi kustom dan daftar ke TFLite Interpreter
Implementasikan operasi gabungan Anda sebagai Operasi kustom TFLite - lihat petunjuk ini.
Perlu diperhatikan bahwa nama untuk mendaftarkan operasi harus sama dengan nama
yang ditetapkan dalam atribut name dalam tanda tangan yang mengimplementasikan.
Contoh op 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 komposit ke operasi gabungan (Lanjutan)
Keseluruhan arsitektur untuk mengonversi operasi komposit TensorFlow menjadi Operasi gabungan LiteRT adalah sebagai berikut:
Menggabungkan operasi gabungan dalam tf.function
Dalam kode sumber model TensorFlow, identifikasi dan pisahkan komposit operasi ke dalam tf.function dengan experimental_implements anotasi fungsi. Lihat contoh pencarian penyematan. Tujuan mendefinisikan antarmuka dan argumennya harus digunakan untuk mengimplementasikan logika konversi.
Menulis kode konversi
Kode konversi ditulis sesuai antarmuka fungsi dengan tag
anotasi implements. Lihat contoh fusi untuk penyematan
pencarian. Secara konseptual, kode konversi menggantikan
implementasi antarmuka ini
dengan antarmuka yang menyatu.
Dalam penerusan fungsi komposit persiapan, plugin di bagian konversi dari kode Anda.
Dalam penggunaan yang lebih maju, dimungkinkan untuk menerapkan transformasi kompleks dari operand operasi komposit untuk memperoleh operand dari gabungan operasi. Lihat Keras LSTM. kode konversi sebagai contoh.
Konversikan ke LiteRT
Gunakan TFLiteConverter.from_saved_model API untuk dikonversi ke LiteRT.
Di balik layar
Kami sekarang menjelaskan detail tingkat tinggi dari keseluruhan desain dalam mengonversi data menjadi gabungan di LiteRT.
Operasi penulisan di TensorFlow
Penggunaan tf.function dengan experimental_implements memungkinkan pengguna untuk secara eksplisit menulis operasi baru menggunakan Operasi primitif TensorFlow dan menentukan antarmuka yang dihasilkan akan diterapkan oleh operasi gabungan. Hal ini sangat berguna karena menyediakan:
- Batas yang jelas untuk operasi gabungan dalam lapisan Grafik TensorFlow.
- Tentukan secara eksplisit antarmuka yang diimplementasikan operasi ini. Tujuan argumen tf.function sesuai dengan argumen antarmuka ini.
Sebagai contoh, mari kita pertimbangkan operasi gabungan yang ditentukan untuk diimplementasikan pencarian embedding. Hal ini dipetakan ke operasi gabungan di LiteRT.
@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 gabungan melalui tf.function sebagai diilustrasikan di atas, membangun infrastruktur umum menjadi mungkin mengidentifikasi dan mengonversi operasi tersebut menjadi operasi LiteRT gabungan.
Memperluas pengonversi LiteRT
Konverter LiteRT yang dirilis awal tahun ini hanya mendukung mengimpor model TensorFlow sebagai grafik dengan semua variabel diganti dengan nilai konstanta yang sesuai. Ini tidak berfungsi untuk fusi operasi karena grafik tersebut memiliki semua fungsi {i>inline<i} sehingga variabel dapat diubah menjadi konstanta.
Untuk memanfaatkan model
tf.function dengan
experimental_implements selama proses konversi, fungsi
harus dipertahankan hingga
proses konversi nanti.
Oleh karena itu, kami menerapkan alur kerja baru untuk mengimpor dan mengonversi TensorFlow model di konverter untuk mendukung kasus penggunaan fusi operasi gabungan. Secara khusus, fitur baru yang ditambahkan adalah:
- Mengimpor model tersimpan TensorFlow ke MLIR
- operasi gabungan sekring
- analisis mutabilitas variabel
- bekukan semua variabel hanya baca
Hal ini memungkinkan kita untuk melakukan fusi operasi menggunakan fungsi yang mewakili operasi gabungan sebelum fungsi inline dan pembekuan variabel.
Mengimplementasikan fusi operasi
Mari kita lihat penerusan fusi operasi secara lebih mendetail. Kartu ini melakukan berikut ini:
- Melakukan loop semua fungsi dalam modul MLIR.
- Jika suatu fungsi memiliki atribut tf._ Implements, berdasarkan atribut tersebut nilai, memanggil utilitas fusi operasi yang sesuai.
- Utilitas fusi operasi beroperasi pada operand fungsi dan (yang berfungsi sebagai antarmuka untuk konversi) dan menggantikan isi fungsi dengan isi fungsi yang setara yang berisi operasi yang menyatu.
- Biasanya, isi yang diganti akan berisi operasi selain operasi yang menyatu. Ini berkaitan dengan beberapa transformasi statis pada operand untuk mendapatkan operand dari operasi gabungan. Karena semua komputasi ini bisa selalu dilipat, maka tidak akan ada dalam {i>flatbuffer<i} 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 sebuah gabungan operasi di LiteRT dengan 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());
}