Thanks for tuning in to Google I/O. View all sessions on demandWatch on demand

Operator khusus

Karena pustaka operator bawaan TensorFlow Lite hanya mendukung operator TensorFlow dalam jumlah terbatas, tidak semua model dapat dikonversi. Untuk detailnya, lihat kompatibilitas operator .

Untuk mengizinkan konversi, pengguna dapat menyediakan implementasi kustom mereka sendiri dari operator TensorFlow yang tidak didukung di TensorFlow Lite, yang dikenal sebagai operator kustom. Jika sebaliknya, Anda ingin menggabungkan serangkaian operator TensorFlow yang tidak didukung (atau didukung) menjadi satu operator kustom yang dioptimalkan dengan fusi, lihat operator fusing .

Menggunakan operator kustom terdiri dari empat langkah.

Mari kita telusuri contoh end-to-end menjalankan model dengan operator kustom tf.atan (dinamai sebagai Atan , lihat #create_a_tensorflow_model) yang didukung di TensorFlow, tetapi tidak didukung di TensorFlow Lite.

Operator Teks TensorFlow adalah contoh operator kustom. Lihat tutorial Konversi Teks TF ke TF Lite untuk contoh kode.

Contoh: Operator Custom Atan

Mari kita telusuri contoh mendukung operator TensorFlow yang tidak dimiliki TensorFlow Lite. Asumsikan kita menggunakan operator Atan dan kita sedang membangun model yang sangat sederhana untuk sebuah fungsi y = atan(x + offset) , di mana offset dapat dilatih.

Buat Model TensorFlow

Cuplikan kode berikut melatih model TensorFlow sederhana. Model ini hanya berisi operator khusus bernama Atan , yang merupakan fungsi y = atan(x + offset) , di mana offset dapat dilatih.

import tensorflow as tf

# Define training dataset and variables
x = [-8, 0.5, 2, 2.2, 201]
y = [-1.4288993, 0.98279375, 1.2490457, 1.2679114, 1.5658458]
offset = tf.Variable(0.0)

# Define a simple model which just contains a custom operator named `Atan`
@tf.function(input_signature=[tf.TensorSpec.from_tensor(tf.constant(x))])
def atan(x):
  return tf.atan(x + offset, name="Atan")

# Train model
optimizer = tf.optimizers.Adam(0.01)
def train(x, y):
    with tf.GradientTape() as t:
      predicted_y = atan(x)
      loss = tf.reduce_sum(tf.square(predicted_y - y))
    grads = t.gradient(loss, [offset])
    optimizer.apply_gradients(zip(grads, [offset]))

for i in range(1000):
    train(x, y)

print("The actual offset is: 1.0")
print("The predicted offset is:", offset.numpy())

The actual offset is: 1.0
The predicted offset is: 0.99999905

Pada titik ini, jika Anda mencoba membuat model TensorFlow Lite dengan bendera konverter default, Anda akan mendapatkan pesan kesalahan berikut:

Error:
error: 'tf.Atan' op is neither a custom op nor a flex op.

Konversikan ke Model TensorFlow Lite

Buat model TensorFlow Lite dengan operator khusus, dengan menyetel atribut konverter allow_custom_ops seperti yang ditunjukkan di bawah ini:

converter = tf.lite.TFLiteConverter.from_concrete_functions([atan.get_concrete_function()], atan)
converter.allow_custom_ops = True
tflite_model = converter.convert()

Pada titik ini, jika Anda menjalankannya dengan juru bahasa default menggunakan perintah seperti berikut:

interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()

Anda masih akan mendapatkan kesalahan:

Encountered unresolved custom op: Atan.

Buat dan daftarkan operator.

Semua operator TensorFlow Lite (kustom dan bawaan) ditentukan menggunakan antarmuka pure-C sederhana yang terdiri dari empat fungsi:

typedef struct {
  void* (*init)(TfLiteContext* context, const char* buffer, size_t length);
  void (*free)(TfLiteContext* context, void* buffer);
  TfLiteStatus (*prepare)(TfLiteContext* context, TfLiteNode* node);
  TfLiteStatus (*invoke)(TfLiteContext* context, TfLiteNode* node);
} TfLiteRegistration;

Lihat common.h untuk detail tentang TfLiteContext dan TfLiteNode . Yang pertama menyediakan fasilitas pelaporan kesalahan dan akses ke objek global, termasuk semua tensor. Yang terakhir memungkinkan implementasi untuk mengakses input dan output mereka.

Saat interpreter memuat model, ia memanggil init() satu kali untuk setiap node dalam grafik. init() yang diberikan akan dipanggil lebih dari sekali jika op digunakan berkali-kali dalam grafik. Untuk operasi khusus, buffer konfigurasi akan disediakan, berisi flexbuffer yang memetakan nama parameter ke nilainya. Buffer kosong untuk operasi bawaan karena juru bahasa telah mem-parsing parameter operasi. Implementasi kernel yang memerlukan status harus diinisialisasi di sini dan mentransfer kepemilikan ke pemanggil. Untuk setiap panggilan init() , akan ada panggilan terkait ke free() , yang memungkinkan implementasi membuang buffer yang mungkin telah dialokasikan di init() .

Setiap kali tensor input diubah ukurannya, juru bahasa akan menelusuri grafik yang memberitahukan implementasi perubahan tersebut. Ini memberi mereka kesempatan untuk mengubah ukuran buffer internal mereka, memeriksa validitas bentuk dan jenis masukan, dan menghitung ulang bentuk keluaran. Ini semua dilakukan melalui prepare() , dan implementasi dapat mengakses statusnya menggunakan node->user_data .

Terakhir, setiap kali inferensi berjalan, penafsir melintasi pemanggilan grafik invoke() , dan di sini juga status tersedia sebagai node->user_data .

Operasi khusus dapat diimplementasikan dengan cara yang persis sama seperti operasi bawaan, dengan mendefinisikan keempat fungsi tersebut dan fungsi pendaftaran global yang biasanya terlihat seperti ini:

namespace tflite {
namespace ops {
namespace custom {
  TfLiteRegistration* Register_MY_CUSTOM_OP() {
    static TfLiteRegistration r = {my_custom_op::Init,
                                   my_custom_op::Free,
                                   my_custom_op::Prepare,
                                   my_custom_op::Eval};
    return &r;
  }
}  // namespace custom
}  // namespace ops
}  // namespace tflite

Perhatikan bahwa pendaftaran tidak otomatis dan panggilan eksplisit ke Register_MY_CUSTOM_OP harus dilakukan. Sementara BuiltinOpResolver standar (tersedia dari :builtin_ops target) menangani pendaftaran builtin, operasi khusus harus dikumpulkan di perpustakaan khusus yang terpisah.

Menentukan kernel di runtime TensorFlow Lite

Yang perlu kita lakukan untuk menggunakan op di TensorFlow Lite adalah menentukan dua fungsi ( Prepare dan Eval ), dan membuat TfLiteRegistration :

TfLiteStatus AtanPrepare(TfLiteContext* context, TfLiteNode* node) {
  using namespace tflite;
  TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);

  const TfLiteTensor* input = GetInput(context, node, 0);
  TfLiteTensor* output = GetOutput(context, node, 0);

  int num_dims = NumDimensions(input);

  TfLiteIntArray* output_size = TfLiteIntArrayCreate(num_dims);
  for (int i=0; i<num_dims; ++i) {
    output_size->data[i] = input->dims->data[i];
  }

  return context->ResizeTensor(context, output, output_size);
}

TfLiteStatus AtanEval(TfLiteContext* context, TfLiteNode* node) {
  using namespace tflite;
  const TfLiteTensor* input = GetInput(context, node, 0);
  TfLiteTensor* output = GetOutput(context, node, 0);

  float* input_data = GetTensorData<float>(input);
  float* output_data = GetTensorData<float>(output);

  size_t count = 1;
  int num_dims = NumDimensions(input);
  for (int i = 0; i < num_dims; ++i) {
    count *= input->dims->data[i];
  }

  for (size_t i=0; i<count; ++i) {
    output_data[i] = atan(input_data[i]);
  }
  return kTfLiteOk;
}

TfLiteRegistration* Register_ATAN() {
  static TfLiteRegistration r = {nullptr, nullptr, AtanPrepare, AtanEval};
  return &r;
}

Saat menginisialisasi OpResolver , tambahkan op kustom ke dalam resolver (lihat contoh di bawah). Ini akan mendaftarkan operator dengan Tensorflow Lite sehingga TensorFlow Lite dapat menggunakan implementasi baru. Perhatikan bahwa dua argumen terakhir di TfLiteRegistration sesuai dengan fungsi AtanPrepare dan AtanEval yang Anda tentukan untuk op kustom. Jika Anda menggunakan fungsi AtanInit dan AtanFree untuk menginisialisasi variabel yang digunakan dalam op dan mengosongkan ruang, masing-masing, keduanya akan ditambahkan ke dua argumen pertama TfLiteRegistration ; argumen tersebut diatur ke nullptr dalam contoh ini.

Daftarkan operator dengan pustaka kernel

Sekarang kita perlu mendaftarkan operator dengan pustaka kernel. Ini dilakukan dengan OpResolver . Di belakang layar, interpreter akan memuat pustaka kernel yang akan ditugaskan untuk mengeksekusi setiap operator dalam model. Meskipun pustaka default hanya berisi kernel bawaan, dimungkinkan untuk mengganti/menambahnya dengan operator op pustaka khusus.

Kelas OpResolver , yang menerjemahkan kode dan nama operator menjadi kode sebenarnya, didefinisikan seperti ini:

class OpResolver {
  virtual TfLiteRegistration* FindOp(tflite::BuiltinOperator op) const = 0;
  virtual TfLiteRegistration* FindOp(const char* op) const = 0;
  virtual void AddBuiltin(tflite::BuiltinOperator op, TfLiteRegistration* registration) = 0;
  virtual void AddCustom(const char* op, TfLiteRegistration* registration) = 0;
};

Penggunaan reguler mengharuskan Anda menggunakan BuiltinOpResolver dan menulis:

tflite::ops::builtin::BuiltinOpResolver resolver;

Untuk menambahkan operasi khusus yang dibuat di atas, Anda memanggil AddOp (sebelum meneruskan penyelesai ke InterpreterBuilder ):

resolver.AddCustom("Atan", Register_ATAN());

Jika kumpulan operasi bawaan dianggap terlalu besar, OpResolver baru dapat dibuat dengan kode berdasarkan subset operasi tertentu, mungkin hanya yang terdapat dalam model tertentu. Ini setara dengan pendaftaran selektif TensorFlow (dan versi sederhananya tersedia di direktori tools ).

Jika Anda ingin menentukan operator khusus di Java, saat ini Anda perlu membuat lapisan JNI khusus Anda sendiri dan mengompilasi AAR Anda sendiri dalam kode jni ini . Demikian pula, jika Anda ingin menentukan operator ini tersedia di Python, Anda dapat menempatkan pendaftaran Anda di kode pembungkus Python .

Perhatikan bahwa proses serupa seperti di atas dapat diikuti untuk mendukung serangkaian operasi, bukan satu operator. Cukup tambahkan operator AddCustom sebanyak yang Anda butuhkan. Selain itu, BuiltinOpResolver juga memungkinkan Anda mengganti implementasi bawaan dengan menggunakan AddBuiltin .

Uji dan buat profil operator Anda

Untuk membuat profil operasi Anda dengan alat tolok ukur TensorFlow Lite, Anda dapat menggunakan alat model tolok ukur untuk TensorFlow Lite. Untuk tujuan pengujian, Anda dapat membuat build lokal TensorFlow Lite mengetahui operasi kustom Anda dengan menambahkan panggilan AddCustom yang sesuai (seperti yang ditampilkan di atas) ke register.cc

Praktik terbaik

  1. Optimalkan alokasi memori dan de-alokasi dengan hati-hati. Mengalokasikan memori di Prepare lebih efisien daripada di Invoke , dan mengalokasikan memori sebelum loop lebih baik daripada di setiap iterasi. Gunakan data tensor sementara daripada melakukan mallocing sendiri (lihat item 2). Gunakan petunjuk/referensi alih-alih menyalin sebanyak mungkin.

  2. Jika struktur data akan bertahan selama keseluruhan operasi, kami menyarankan untuk mengalokasikan memori terlebih dahulu menggunakan tensor sementara. Anda mungkin perlu menggunakan struct OpData untuk mereferensikan indeks tensor di fungsi lain. Lihat contoh di kernel untuk konvolusi . Cuplikan kode contoh ada di bawah

    auto* op_data = reinterpret_cast<OpData*>(node->user_data);
TfLiteIntArrayFree(node->temporaries);
node->temporaries = TfLiteIntArrayCreate(1);
node->temporaries->data[0] = op_data->temp_tensor_index;
TfLiteTensor* temp_tensor = &context->tensors[op_data->temp_tensor_index];
temp_tensor->type =  kTfLiteFloat32;
temp_tensor->allocation_type = kTfLiteArenaRw;

  3. Jika tidak menghabiskan terlalu banyak memori yang terbuang percuma, lebih baik menggunakan array ukuran tetap statis (atau std::vector yang dialokasikan sebelumnya di Resize ) daripada menggunakan std::vector yang dialokasikan secara dinamis setiap iterasi eksekusi.

  4. Hindari pembuatan contoh wadah pustaka standar yang belum ada, karena memengaruhi ukuran biner. Misalnya, jika Anda memerlukan std::map dalam operasi Anda yang tidak ada di kernel lain, menggunakan std::vector dengan pemetaan pengindeksan langsung dapat berfungsi sambil menjaga ukuran biner tetap kecil. Lihat apa yang digunakan kernel lain untuk mendapatkan wawasan (atau bertanya).

  5. Periksa penunjuk ke memori yang dikembalikan oleh malloc . Jika penunjuk ini nullptr , tidak ada operasi yang harus dilakukan menggunakan penunjuk itu. Jika Anda malloc dalam suatu fungsi dan memiliki kesalahan keluar, batalkan alokasi memori sebelum Anda keluar.

  6. Gunakan TF_LITE_ENSURE(context, condition) untuk memeriksa kondisi tertentu. Kode Anda tidak boleh membiarkan memori menggantung saat TF_LITE_ENSURE digunakan, yaitu, makro ini harus digunakan sebelum sumber daya yang dialokasikan akan bocor.