Toán tử tùy chỉnh

Sử dụng bộ sưu tập để sắp xếp ngăn nắp các trang Lưu và phân loại nội dung dựa trên lựa chọn ưu tiên của bạn.

Vì thư viện toán tử nội trang TensorFlow Lite chỉ hỗ trợ một số giới hạn toán tử TensorFlow, 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ọ về 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 các 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 hợp nhất được tối ưu hóa duy nhất, hãy tham khảo kết hợp 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ụ từ đầu đến cuối về việc chạy một mô hình với toán tử tùy chỉnh tf.sin (có tên là Sin , hãy 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ử TensorFlow Text là một ví dụ về toán tử tùy chỉnh. Xem hướng dẫn Chuyển đổi văn bản TF sang TF Lite để biết ví dụ về mã.

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

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

Tạo 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 một toán tử tùy chỉnh có tên là Sin , là một hàm y = sin(x + offset) , trong đó offset có thể đào tạo được.

import tensorflow as tf

# Define training dataset and variables
x = [-8, 0.5, 2, 2.2, 201]
y = [-0.6569866 ,  0.99749499,  0.14112001, -0.05837414,  0.80641841]
offset = tf.Variable(0.0)

# Define a simple model which just contains a custom operator named `Sin`
@tf.function
def sin(x):
  return tf.sin(x + offset, name="Sin")

# Train model
optimizer = tf.optimizers.Adam(0.01)
def train(x, y):
    with tf.GradientTape() as t:
      predicted_y = sin(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: 1.0000001

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

Error:
Some of the operators in the model are not supported by the standard TensorFlow
Lite runtime...... Here is
a list of operators for which you will need custom implementations: Sin.

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 chuyển đổi allow_custom_ops như được hiển thị bên dưới:

converter = tf.lite.TFLiteConverter.from_concrete_functions([sin.get_concrete_function(x)], sin)
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ạn sẽ nhận được thông báo lỗi sau:

Error:
Didn't find custom operator for name 'Sin'
Registration failed.

Tạo và đăng ký nhà điều hành.

Tất cả các toán tử TensorFlow Lite (cả tùy chỉnh và nội trang) đều được xác định bằng giao diện pure-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ề TfLiteContextTfLiteNode . Trước đây 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 tensor. Cái sau cho phép các triển khai truy cập đầ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 thông số với giá trị của chúng. Bộ đệm trống cho các hoạt động nội trang vì trình thông dịch đã phân tích cú pháp các tham số hoạt động. Các triển khai hạt nhân 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 các triển khai loại bỏ bộ đệm mà chúng có thể đã cấp phát trong init() .

Bất cứ khi nào các bộ căng đầ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 giúp họ có 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 hình dạng và kiểu đầ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à cá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 suy luận chạy, trình thông dịch 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 .

Hoạt động tùy chỉnh có thể được triển khai theo cách giống hệt như hoạt động nội trang, bằng cách xác định bốn chức năng đó và một chức năng đăng ký toàn cục thường trô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 một cuộc gọi rõ ràng tới Register_MY_CUSTOM_OP . Trong khi BuiltinOpResolver tiêu chuẩn (có sẵn từ đích :builtin_ops ) quản lý việc đăng ký nội trang, các hoạt động 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.

Định nghĩa hạt nhân 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 ( PrepareTfLiteRegistration Eval

TfLiteStatus SinPrepare(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 SinEval(TfLiteContext* context, TfLiteNode* node) {
  using namespace tflite;
  const TfLiteTensor* input = GetInput(context, node,0);
  TfLiteTensor* output = GetOutput(context, node,0);

  float* input_data = input->data.f;
  float* output_data = output->data.f;

  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] = sin(input_data[i]);
  }
  return kTfLiteOk;
}

TfLiteRegistration* Register_SIN() {
  static TfLiteRegistration r = {nullptr, nullptr, SinPrepare, SinEval};
  return &r;
}

Khi khởi tạo OpResolver , hãy thêm tùy chọn 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ý nhà điều hành 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 SinPrepareSinEval mà bạn đã xác định cho tùy chọn tùy chỉnh. Nếu bạn đã sử dụng các SinInitSinFree để 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 ; các đố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 OpResolver . Phía sau, trình thông dịch sẽ tải một thư viện các 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 nhân nội trang, nhưng có thể thay thế / bổ sung nó bằng các toán tử op thư viện tùy chỉnh.

Lớp OpResolver , dịch các mã và tên của toán tử thành mã thực, đượ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 đã tạo ở trên, bạn gọi AddOp (trước khi bạn chuyển trình phân giải cho InterpreterBuilder ):

resolver.AddCustom("Sin", Register_SIN());

Nếu tập hợp các hoạt động tích hợp được coi là quá lớn, một OpResolver mới có thể được tạo mã dựa trên một tập hợp con các hoạt động nhất định, có thể chỉ những hoạt động có trong một mô hình nhất định. Đây là 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 xác định các toán tử tùy chỉnh của mình trong Java, bạn hiện cần phải 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 quy trình tương tự như trên có thể được tuân theo để hỗ trợ một tập hợp các hoạt động thay vì một toán tử duy nhất. Chỉ cần thêm nhiều toán tử AddCustom nếu bạn cần. Ngoài ra, BuiltinOpResolver cũng cho phép bạn ghi đè triển khai các nội trang bằng cách sử dụng AddBuiltin .

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

Để lập hồ sơ lựa chọn của bạn với 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. Với mục đích thử nghiệm, bạn có thể làm cho phiên bản TensorFlow Lite cục bộ của mình biết về tùy chọn 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) để đăng ký.cc

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

  1. Tối ưu hóa cấp phát bộ nhớ và hủy cấp phát một cách thận trọng. Phân bổ bộ nhớ trong Prepare 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. Sử dụng dữ liệu căng tạm thời thay vì tự làm hỏng (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 cấu trúc dữ liệu vẫn tồn tại trong toàn bộ hoạt động, chúng tôi khuyên bạn nên cấp phát trước bộ nhớ bằng cách sử dụng bộ căng 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 hàm khác. Xem ví dụ trong nhân để biết tích chập . Dưới đây là một đoạn mã mẫu

    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ớ lãng phí, hãy thích sử dụng mảng kích thước cố định tĩnh (hoặc std::vector được cấp phát trước trong Resize ) hơn là sử dụng std::vector được cấp phát độ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 mà không tồn tại trong các hạt nhân khác, 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ì các hạt nhân khác sử dụng để có được cái nhìn sâu sắc (hoặc hỏi).

  5. Kiểm tra con trỏ đến bộ nhớ do malloc trả về. Nếu con trỏ này là nullptr , không có thao tác nào được thực hiện bằng con trỏ đó. Nếu bạn nhập malloc một hàm và thoát ra lỗi, hãy phân bổ bộ nhớ trước khi 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 để bộ nhớ bị treo khi TF_LITE_ENSURE được sử dụng, tức là, các macro này phải được sử dụng trước khi bất kỳ tài nguyên nào được cấp phát sẽ bị rò rỉ.