Toán tử tùy chỉnh

Do thư viện toán tử dựng sẵn TensorFlow Lite chỉ hỗ trợ một số lượng hạn chế toán tử TensorFlow nên không phải mọi mô hình đều có thể chuyển đổi. Để biết chi tiết, hãy tham khảo khả năng tương thích của nhà điều hành .

Để cho phép chuyển đổi, người dùng có thể cung cấp triển khai tùy chỉnh của riêng họ cho toán tử TensorFlow không được hỗ trợ trong TensorFlow Lite, được gọi là toán tử tùy chỉnh. Thay vào đó, nếu bạn muốn kết hợp một loạt toán tử TensorFlow không được hỗ trợ (hoặc được hỗ trợ) thành một toán tử tùy chỉnh được tối ưu hóa hợp nhất duy nhất, hãy tham khảo hợp nhất toán tử .

Sử dụng toán tử tùy chỉnh bao gồm bốn bước.

Hãy xem qua một ví dụ toàn diện về cách chạy một mô hình với toán tử tùy chỉnh tf.atan (được đặt tên là Atan , tham khảo #create_a_tensorflow_model) được hỗ trợ trong TensorFlow nhưng không được hỗ trợ trong TensorFlow Lite.

Toán tử Văn bản TensorFlow là một ví dụ về toán tử tùy chỉnh. Xem hướng dẫn Chuyển văn bản TF sang TF Lite để biết ví dụ về mã.

Ví dụ: Toán tử Atan tùy chỉnh

Hãy xem qua một ví dụ về hỗ trợ toán tử TensorFlow mà TensorFlow Lite không có. Giả sử chúng ta đang sử dụng toán tử Atan và chúng ta đang xây dựng một mô hình rất đơn giản cho hàm y = atan(x + offset) , trong đó offset có thể huấn luyện được.

Tạo một mô hình TensorFlow

Đoạn mã sau đào tạo một mô hình TensorFlow đơn giản. Mô hình này chỉ chứa toán tử tùy chỉnh có tên Atan , là một hàm y = atan(x + offset) , trong đó offset có thể huấn luyện được.

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

Tại thời điểm này, nếu cố gắng tạo mô hình TensorFlow Lite với các cờ chuyển đổi mặc định, bạn sẽ nhận được thông báo lỗi sau:

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

Chuyển đổi sang Mô hình TensorFlow Lite

Tạo mô hình TensorFlow Lite với các toán tử tùy chỉnh, bằng cách đặt thuộc tính trình chuyển đổi allow_custom_ops như hiển thị bên dưới:

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

Tại thời điểm này, nếu bạn chạy nó với trình thông dịch mặc định bằng cách sử dụng các lệnh như sau:

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

Bạn vẫn sẽ gặp lỗi:

Encountered unresolved custom op: Atan.

Tạo và đăng ký toán tử.

Tất cả các toán tử TensorFlow Lite (cả tùy chỉnh và dựng sẵn) đều được xác định bằng giao diện thuần C đơn giản bao gồm bốn chức năng:

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;

Tham khảo common.h để biết chi tiết về TfLiteContext và TfLiteNode . Cái trước cung cấp các phương tiện báo cáo lỗi và quyền truy cập vào các đối tượng toàn cầu, bao gồm tất cả các tenxơ. Cái sau cho phép việc triển khai truy cập vào đầu vào và đầu ra của chúng.

Khi trình thông dịch tải một mô hình, nó sẽ gọi init() một lần cho mỗi nút trong biểu đồ. Một init() đã cho sẽ được gọi nhiều lần nếu op được sử dụng nhiều lần trong biểu đồ. Đối với các hoạt động tùy chỉnh, bộ đệm cấu hình sẽ được cung cấp, chứa bộ đệm linh hoạt ánh xạ tên tham số với giá trị của chúng. Bộ đệm trống cho các op dựng sẵn vì trình thông dịch đã phân tích cú pháp các tham số op. Việc triển khai kernel yêu cầu trạng thái nên khởi tạo nó ở đây và chuyển quyền sở hữu cho người gọi. Đối với mỗi lệnh gọi init() , sẽ có một lệnh gọi tương ứng tới free() , cho phép việc triển khai loại bỏ bộ đệm mà chúng có thể đã phân bổ trong init() .

Bất cứ khi nào các tenxơ đầu vào được thay đổi kích thước, trình thông dịch sẽ đi qua biểu đồ thông báo việc triển khai thay đổi. Điều này mang lại cho họ cơ hội thay đổi kích thước bộ đệm bên trong, kiểm tra tính hợp lệ của các loại và hình dạng đầu vào cũng như tính toán lại các hình dạng đầu ra. Tất cả điều này được thực hiện thông qua prepare() và việc triển khai có thể truy cập trạng thái của chúng bằng cách sử dụng node->user_data .

Cuối cùng, mỗi khi chạy suy luận, trình thông dịch sẽ duyệt qua biểu đồ gọi invoke() và ở đây trạng thái cũng có sẵn dưới dạng node->user_data .

Các op tùy chỉnh có thể được triển khai theo cách chính xác giống như các op dựng sẵn, bằng cách xác định bốn hàm đó và một hàm đăng ký toàn cầu thường giống như sau:

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

Lưu ý rằng đăng ký không tự động và phải thực hiện cuộc gọi rõ ràng tới Register_MY_CUSTOM_OP . Mặc dù BuiltinOpResolver tiêu chuẩn (có sẵn từ mục tiêu :builtin_ops ) đảm nhận việc đăng ký các nội trang, các op tùy chỉnh sẽ phải được thu thập trong các thư viện tùy chỉnh riêng biệt.

Xác định kernel trong thời gian chạy TensorFlow Lite

Tất cả những gì chúng ta cần làm để sử dụng op trong TensorFlow Lite là xác định hai chức năng ( Prepare và Eval ) và xây dựng một 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;
}

Khi khởi tạo OpResolver , hãy thêm op tùy chỉnh vào trình phân giải (xem ví dụ bên dưới). Thao tác này sẽ đăng ký toán tử với Tensorflow Lite để TensorFlow Lite có thể sử dụng triển khai mới. Lưu ý rằng hai đối số cuối cùng trong TfLiteRegistration tương ứng với các hàm AtanPrepare và AtanEval mà bạn đã xác định cho op tùy chỉnh. Nếu bạn đã sử dụng các hàm AtanInit và AtanFree để khởi tạo các biến được sử dụng trong op và để giải phóng dung lượng, thì chúng sẽ được thêm vào hai đối số đầu tiên của TfLiteRegistration ; những đối số đó được đặt thành nullptr trong ví dụ này.

Đăng ký toán tử với thư viện hạt nhân

Bây giờ chúng ta cần đăng ký toán tử với thư viện hạt nhân. Điều này được thực hiện với một OpResolver . Đằng sau hậu trường, trình thông dịch sẽ tải một thư viện hạt nhân sẽ được chỉ định để thực thi từng toán tử trong mô hình. Mặc dù thư viện mặc định chỉ chứa các kernel dựng sẵn, nhưng có thể thay thế/tăng cường nó bằng một toán tử op thư viện tùy chỉnh.

Lớp OpResolver , lớp dịch mã và tên toán tử thành mã thực tế, được định nghĩa như sau:

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;
};

Việc sử dụng thông thường yêu cầu bạn sử dụng BuiltinOpResolver và viết:

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

Để thêm op tùy chỉnh được tạo ở trên, bạn gọi AddOp (trước khi bạn chuyển trình phân giải tới InterpreterBuilder ):

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

Nếu tập hợp các ops dựng sẵn được coi là quá lớn, thì một OpResolver mới có thể được tạo mã dựa trên một tập hợp con ops nhất định, có thể chỉ những ops có trong một mô hình nhất định. Điều này tương đương với đăng ký chọn lọc của TensorFlow (và một phiên bản đơn giản của nó có sẵn trong thư mục tools ).

Nếu bạn muốn định nghĩa các toán tử tùy chỉnh của mình trong Java, hiện tại bạn cần xây dựng lớp JNI tùy chỉnh của riêng mình và biên dịch AAR của riêng bạn trong mã jni này . Tương tự, nếu bạn muốn xác định các toán tử này có sẵn trong Python, bạn có thể đặt đăng ký của mình trong mã trình bao bọc Python .

Lưu ý rằng có thể tuân theo quy trình tương tự như trên để hỗ trợ một tập hợp các thao tác thay vì một toán tử đơn lẻ. Chỉ cần thêm bao nhiêu toán tử AddCustom mà bạn cần. Ngoài ra, BuiltinOpResolver cũng cho phép bạn ghi đè việc triển khai nội trang bằng cách sử dụng AddBuiltin .

Kiểm tra và lập hồ sơ nhà điều hành của bạn

Để lập hồ sơ hoạt động của bạn bằng công cụ điểm chuẩn TensorFlow Lite, bạn có thể sử dụng công cụ mô hình điểm chuẩn cho TensorFlow Lite. Đối với mục đích thử nghiệm, bạn có thể làm cho bản dựng TensorFlow Lite cục bộ của mình biết về hoạt động tùy chỉnh của bạn bằng cách thêm lệnh gọi AddCustom thích hợp (như hiển thị ở trên) vào register.cc

Thực hành tốt nhất

  1. Tối ưu hóa phân bổ bộ nhớ và phân bổ lại một cách thận trọng. Cấp phát bộ nhớ trong Prepare bị hiệu quả hơn trong Invoke và cấp phát bộ nhớ trước một vòng lặp tốt hơn trong mỗi lần lặp lại. Sử dụng dữ liệu tenxơ tạm thời thay vì tự đánh lừa bản thân (xem mục 2). Sử dụng con trỏ/tham chiếu thay vì sao chép càng nhiều càng tốt.

  2. Nếu một cấu trúc dữ liệu sẽ tồn tại trong toàn bộ hoạt động, chúng tôi khuyên bạn nên phân bổ trước bộ nhớ bằng cách sử dụng các tenxơ tạm thời. Bạn có thể cần sử dụng cấu trúc OpData để tham chiếu các chỉ số tensor trong các chức năng khác. Xem ví dụ trong hạt nhân để biết tích chập . Một đoạn mã mẫu dưới đây

    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. Nếu nó không tốn quá nhiều bộ nhớ bị lãng phí, hãy ưu tiên sử dụng một mảng có kích thước cố định tĩnh (hoặc một std::vector được phân bổ trước trong Resize ) hơn là sử dụng một std::vector được phân bổ động mỗi lần lặp lại thực thi.

  4. Tránh khởi tạo các mẫu vùng chứa thư viện tiêu chuẩn chưa tồn tại vì chúng ảnh hưởng đến kích thước nhị phân. Ví dụ: nếu bạn cần một std::map trong hoạt động của mình không tồn tại trong các hạt nhân khác, thì việc sử dụng một std::vector với ánh xạ lập chỉ mục trực tiếp có thể hoạt động trong khi vẫn giữ kích thước nhị phân nhỏ. Xem những gì hạt nhân khác sử dụng để đạt được cái nhìn sâu sắc (hoặc hỏi).

  5. Kiểm tra con trỏ tới bộ nhớ được trả về bởi malloc . Nếu con trỏ này là nullptr , thì không có thao tác nào được thực hiện bằng con trỏ đó. Nếu bạn malloc trong một chức năng và gặp lỗi thoát, hãy giải phóng bộ nhớ trước khi bạn thoát.

  6. Sử dụng TF_LITE_ENSURE(context, condition) để kiểm tra một điều kiện cụ thể. Mã của bạn không được để treo bộ nhớ khi TF_LITE_ENSURE được sử dụng, nghĩa là các macro này nên được sử dụng trước khi bất kỳ tài nguyên nào được phân bổ sẽ bị rò rỉ.