Toán tử tùy chỉnh

Vì 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 được. Để 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 cách 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 phần toán tử hợp nhất .

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

Chúng ta hãy xem qua một ví dụ chi tiết về việc 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 đổi 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à đ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ô hình TensorFlow

Đoạn mã sau huấn luyện 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 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 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ó bằng trình thông dịch mặc định bằ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ề TfLiteContextTfLiteNode . Cái trước cung cấp phương tiện báo cáo lỗi và quyền truy cập vào các đối tượng chung, bao gồm tất cả các tensor. 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() nhất định 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ố tớ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 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 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 tensor đầu vào được thay đổi kích thước, trình thông dịch sẽ xem 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 hình dạng và loại đầu vào cũng như tính toán lại 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 lần suy luận chạy, 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 là node->user_data .

Các hoạt động tùy chỉnh có thể được triển khai theo cách giống hệt như các hoạt động dựng sẵn, bằng cách xác định bốn chức năng đó và chức năng đăng ký chung thường trông như thế này:

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 việc đăng ký không tự động và phải thực hiện lệnh gọi rõ ràng tới Register_MY_CUSTOM_OP . Trong khi BuiltinOpResolver tiêu chuẩn (có sẵn từ mục tiêu :builtin_ops ) đảm nhiệm việc đăng ký nội dung, 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 hàm ( PrepareEval ) và xây dựng 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 cách 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 AtanPrepareAtanEval mà bạn đã xác định cho op tùy chỉnh. Nếu bạn đã sử dụng các hàm AtanInitAtanFree để khởi tạo các biến được sử dụng trong op và để giải phóng dung lượng tương ứ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 kernel

Bây giờ chúng ta cần đăng ký toán tử với thư viện kernel. Việc này được thực hiện bằng OpResolver . Phía sau, 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 hạt nhân dựng sẵn, nhưng có thể thay thế/tăng cường nó bằng các toán tử op thư viện tùy chỉnh.

Lớp OpResolver , dịch mã toán tử và tên 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 chuyển trình phân giải tới InterpreterBuilder ):

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

Nếu tập hợp các hoạt động dựng sẵn được cho 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 các hoạt động nhất định, có thể chỉ những tập hợp hoạt động 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 muốn xác định các toán tử tùy chỉnh của mình trong Java, thì hiện tại bạ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 vào 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ử duy nhất. Chỉ cần thêm bao nhiêu toán tử AddCustom tùy theo nhu cầu của bạn. Ngoài ra, BuiltinOpResolver còn 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ụ đo đ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 nhận biết được hoạt động tùy chỉnh của mình 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 việc phân bổ và phân bổ bộ nhớ một cách thận trọng. Việc phân bổ bộ nhớ trong Prepare bị hiệu quả hơn trong Invoke và việc phân bổ bộ nhớ trước một vòng lặp sẽ tốt hơn trong mỗi lần lặp. Sử dụng dữ liệu tensor tạm thời thay vì tự điều chỉnh (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 phân bổ trước bộ nhớ bằng cách sử dụng các tensor 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 kernel để biết tích chập . Một đoạn mã mẫu bên dưới

    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 không tốn quá nhiều bộ nhớ, hãy ưu tiên sử dụng mảng có kích thước cố định tĩnh (hoặc std::vector được phân bổ trước trong Resize ) thay vì sử dụng std::vector được phân bổ động mỗi lần lặp 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 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, thì việc sử dụng 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 để hiểu rõ hơn (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 nên thực hiện thao tác nào bằng con trỏ đó. Nếu bạn malloc trong một hàm và gặp lỗi thoát, hãy giải phóng 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 sử dụng TF_LITE_ENSURE , nghĩa 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 phân bổ sẽ bị rò rỉ.