自訂運算子

由於 LiteRT 內建運算子程式庫僅支援有限數量的 TensorFlow 運算子，因此並非所有模型都能轉換。詳情請參閱電信業者相容性。

如要允許轉換，使用者可以在 LiteRT 中提供不受支援的 TensorFlow 運算子自訂實作項目，也就是自訂運算子。如要將一系列不支援 (或支援) 的 TensorFlow 運算子合併為單一融合的最佳化自訂運算子，請參閱運算子融合。

使用自訂運算子需要四個步驟。

我們將逐步說明端對端範例，瞭解如何使用 TensorFlow 支援但 LiteRT 不支援的自訂運算子 tf.atan (命名為 Atan，請參閱「建立 TensorFlow 模型」)。

TensorFlow Text 運算子就是自訂運算子的例子。如需程式碼範例，請參閱「將 TF Text 轉換為 LiteRT」教學課程。

範例：自訂 Atan 運算子

我們來逐步瞭解如何支援 LiteRT 沒有的 TensorFlow 運算子。假設我們使用 Atan 運算子，並為函式 y = atan(x + offset) 建構非常簡單的模型，其中 offset 可訓練。

建立 TensorFlow 模型

下列程式碼片段會訓練簡單的 TensorFlow 模型。這個模型只包含名為 Atan 的自訂運算子，也就是函式 y = atan(x + offset)，其中 offset 可訓練。

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

此時，如果您嘗試使用預設轉換器標記產生 LiteRT 模型，會收到下列錯誤訊息：

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

轉換為 LiteRT 模型

如要使用自訂運算子建立 LiteRT 模型，請設定轉換器屬性 allow_custom_ops，如下所示：

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

此時，如果您使用下列指令，以預設解譯器執行：

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

您仍會收到錯誤訊息：

Encountered unresolved custom op: Atan.

建立及註冊運算子。

#include "third_party/tensorflow/lite/c/c_api.h"
#include "third_party/tensorflow/lite/c/c_api_opaque.h"

LiteRT 自訂運算子是使用簡單的純 C API 定義，其中包含不透明型別 (TfLiteOperator) 和相關函式。

TfLiteOperator 是不透明型別：

typedef struct TfLiteOperator TfLiteOperator;

TfLiteOperator 會儲存運算子的身分和實作方式。(請注意，運算元與運算元不同，運算元會儲存在呼叫運算元的節點的 LiteRT 圖形節點中)。

此類型的執行個體是透過呼叫 TfLiteOperatorCreate 建構而成，並可透過呼叫 TfLiteOperatorDelete 銷毀。

運算子的身分識別是透過建構函式 TfLiteOperatorCreate 的參數設定：

TfLiteOperator*
TfLiteOperatorCreate(
    TfLiteBuiltinOperator builtin_code,  // Normally `TfLiteBuiltinCustom`.
    const char* custom_name,  // The name of the custom op.
    int version  // Normally `1` for the first version of a custom op.
);

運算子實作項目可使用下列簽章定義「方法」。這些方法都是選用，但如要成功評估運算子，運算子實作至少需要定義及設定 (使用設定器函式) Prepare 和 Invoke 方法。

// Initializes the op from serialized data.
void* Init(TfLiteOpaqueContext* context, const char* buffer, size_t length);

// Deallocates the op.
// The pointer `buffer` is the data previously returned by an Init invocation.
void Free(TfLiteOpaqueContext* context, void* buffer);

// Called when the inputs that this node depends on have been resized.
TfLiteStatus Prepare(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node);

// Called when the node is executed. (Should read node inputs and write to
// node outputs).
TfLiteStatus Invoke(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node);

// Retrieves the async kernel.
TfLiteAsyncKernel AsyncKernel(TfLiteOpaqueContext* context,
                              TfLiteOpaqueNode* node);

您在 op 實作中使用的函式名稱 (或 C++ 的命名空間前置字元) 不必與上述程式碼片段中的函式名稱相符，因為 TF Lite 自訂 op API 只會使用這些函式的位址。事實上，我們建議您在匿名命名空間中或以靜態函式形式宣告這些函式。

但建議您在這些函式名稱中加入運算子名稱做為命名空間或前置字元：

C++

namespace my_namespace::my_custom_op {
  void* Init(TfLiteOpaqueContext* context,
             const char* buffer, size_t length) { ... }
  // ... plus definitions of Free, Prepare, and Invoke ...
}

C

void* MyCustomOpInit(TfLiteOpaqueContext* context,
                     const char* buffer, size_t length) { ... }
// ... plus definitions of MyCustomOpFree, MyCustomOpPrepare, and
// MyCustomOpInvoke.

由於這是 C API，因此這些「方法」會在 TfLiteOperator 型別中實作為 C 函式指標，方法是將實作函式的位址傳遞至對應的 setter 函式 TfLiteOperatorSetMethodName：

void TfLiteOperatorSetInit(
    TfLiteOperator* operator,
    void* (*init)(TfLiteOpaqueContext* context, const char* buffer,
                  size_t length));
void TfLiteOperatorSetFree(
    TfLiteOperator* operator,
    void (*free)(TfLiteOpaqueContext* context, void* data));
void TfLiteOperatorSetPrepare(
    TfLiteOperator* operator,
    TfLiteStatus (*prepare)(TfLiteOpaqueContext* context,
                            TfLiteOpaqueNode* node));
void TfLiteOperatorSetInvoke(
    TfLiteOperator* operator,
    TfLiteStatus (*invoke)(TfLiteOpaqueContext* context,
                           TfLiteOpaqueNode* node));
void TfLiteOperatorSetAsyncKernel(
    TfLiteOperator* operator,
    struct TfLiteAsyncKernel* (*async_kernel)(TfLiteOpaqueContext* context,
                                              TfLiteOpaqueNode* node));

如要瞭解 TfLiteContext 和 TfLiteNode 的詳細資訊，請參閱 common.h。TfLiteContext 提供錯誤回報功能，並可存取全域物件，包括所有張量。TfLiteNode 可讓運算子實作項目存取輸入和輸出內容。

解譯器載入模型時，會針對圖中的每個節點呼叫一次 Init() 方法。如果 op 在圖表中多次使用，系統會多次呼叫指定的 Init()。如果是自訂作業，系統會提供設定緩衝區，其中包含將參數名稱對應至值的 flexbuffer。由於解譯器已剖析作業參數，因此內建作業的緩衝區是空的。需要狀態的 Kernel 實作項目應在此處初始化狀態，並將擁有權轉移給呼叫端。每次呼叫 Init() 時，都會有對應的 Free() 呼叫，讓實作項目能處置可能在 Init() 中分配的緩衝區。

每當輸入張量大小調整時，解譯器會通知實作項目變更，並經過圖表。這讓他們有機會調整內部緩衝區大小、檢查輸入形狀和型別的有效性，以及重新計算輸出形狀。這些全都是透過 Prepare() 方法完成，且實作項目可使用 TfLiteOpaqueNodeGetUserData(node) 存取其狀態。

最後，每次執行推論時，解譯器都會遍歷圖表並呼叫 Invoke() 方法，而狀態也會以 TfLiteOpaqueNodeGetUserData(node) 的形式提供。

您可以定義這些「方法」函式，然後定義一個函式，傳回透過呼叫 TfLiteOperatorCreate 和相關設定器方法建構的 TfLiteOperator 例項，即可實作自訂作業：

C++

namespace my_namespace::my_custom_op {
  namespace {
    void* Init(TfLiteOpaqueContext* context,
               const char* buffer, size_t length) { ... }
    void Free(TfLiteOpaqueContext* context, void* buffer) { ... }
    TfLiteStatus Prepare(TfLiteOpaqueContext* context,
                         TfLiteOpaqueNode* node) { ... }
    TfLiteStatus Invoke(TfLiteOpaqueContext* context,
                        TfLiteOpaqueNode* node) {... }
  };

  const TfLiteOperator* MyCustomOperator() {
    // Singleton instance, intentionally never destroyed.
    static const TfLiteOperator* my_custom_op = ()[] {
        TfLiteOperator* r =
            TfLiteOperatorCreate(
                kTfLiteBuiltinCustom, "MyCustomOp", /*version=*/ 1);
        TfLiteOperatorSetInit(r, Init);
        TfLiteOperatorSetFree(r, Free);
        TfLiteOperatorSetPrepare(r, Prepare);
        TfLiteOperatorSetInvoke(r, Eval);
        return r;
      };
    return my_custom_op;
  }
}  // namespace my_namespace

C

static void* MyCustomOpInit(TfLiteOpaqueContext* context, const char* buffer,
                     size_t length) { ... }
static void MyCustomOpFree(TfLiteOpaqueContext* context, void* buffer) { ... }
static TfLiteStatus MyCustomOpPrepare(TfLiteOpaqueContext* context,
                                      TfLiteOpaqueNode* node) { ... }
static TfLiteStatus MyCustomOpInvoke(TfLiteOpaqueContext* context,
                                     TfLiteOpaqueNode* node) {... }

static TfLiteOperator* MyCustomOpCreate() {
  const TfLiteOperator* r =
      TfLiteOperatorCreate(
          kTfLiteBuiltinCustom, "MyCustomOp", /*version=*/ 1);
  TfLiteOperatorSetInit(r, MyCustomOpInit);
  TfLiteOperatorSetFree(r, MyCustomOpFree);
  TfLiteOperatorSetPrepare(r, MyCustomOpPrepare);
  TfLiteOperatorSetInvoke(r, MyCustomOpEval);
  return r;
}

const TfLiteOperator* MyCustomOperator() {
  // Singleton instance, intentionally never destroyed.
  static const TfLiteOperator* my_custom_op = MyCustomOpCreate();
  return my_custom_op;
}

請注意，註冊並非自動進行，您應明確呼叫 MyCustomOperator 函式 (詳情請見下文)。標準 BuiltinOpResolver (可從 :builtin_ops 目標取得) 會負責註冊內建項目，但自訂作業必須收集在個別的自訂程式庫中。

在 LiteRT 執行階段中定義核心

如要在 LiteRT 中使用運算元，我們只需要定義兩個函式 (Prepare 和 Eval)，以及第三個函式來建構 TfLiteOperator：

C++

namespace atan_op {
  namespace {
    TfLiteStatus AtanPrepare(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) {
      TF_LITE_OPAQUE_ENSURE_EQ(context, TfLiteOpaqueNodeNumInputs(node), 1);
      TF_LITE_OPAQUE_ENSURE_EQ(context, TfLiteOpaqueNodeNumOutputs(node), 1);

      const TfLiteOpaqueTensor* input = TfLiteOpaqueNodeGetInput(context, node, 0);
      TfLiteOpaqueTensor* output = TfLiteOpaqueNodeGetOutput(context, node, 0);

      int num_dims = TfLiteOpaqueTensorNumDimensions(input);

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

      return TfLiteOpaqueContextResizeTensor(context, output, output_size);
    }

    TfLiteStatus AtanEval(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) {
      const TfLiteOpaqueTensor* input = TfLiteOpaqueNodeGetInput(context, node, 0);
      TfLiteOpaqueTensor* output = TfLiteOpaqueNodeGetOutput(context, node, 0);

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

      size_t count = 1;
      int num_dims = TfLiteOpaqueTensorNumDimensions(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;
    }
  }  // anonymous namespace

  const TfLiteOperator* AtanOperator() {
    // Singleton instance, intentionally never destroyed.
    static const TfLiteOperator* atan_op = ()[] {
        auto* r = TfLiteOperatorCreate(
            kTfLiteBuiltinCustom, "ATAN", /*version=*/ 1);
        TfLiteOperatorSetPrepare(r, Prepare);
        TfLiteOperatorSetInvoke(r, Eval);
        return r;
      };
    return atan_op;
  }
}  // namespace atan_op

C

static TfLiteStatus AtanPrepare(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) {
  TF_LITE_OPAQUE_ENSURE_EQ(context, TfLiteOpaqueNodeNumInputs(node), 1);
  TF_LITE_OPAQUE_ENSURE_EQ(context, TfLiteOpaqueNodeNumOutputs(node), 1);

  const TfLiteOpaqueTensor* input = TfLiteOpaqueNodeGetInput(context, node, 0);
  TfLiteOpaqueTensor* output = TfLiteOpaqueNodeGetOutput(context, node, 0);

  int num_dims = TfLiteOpaqueTensorNumDimensions(input);

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

  return TfLiteOpaqueContextResizeTensor(context, output, output_size);
}

static TfLiteStatus AtanEval(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) {
  const TfLiteOpaqueTensor* input = TfLiteOpaqueNodeGetInput(context, node, 0);
  TfLiteOpaqueTensor* output = TfLiteOpaqueNodeGetOutput(context, node, 0);

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

  size_t count = 1;
  int num_dims = TfLiteOpaqueTensorNumDimensions(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;
}

static const TfLiteOperator* AtanOpCreate() {
  TfLiteOperator* r = TfLiteOperatorCreate(
          kTfLiteBuiltinCustom, "ATAN", /*version=*/ 1);
  TfLiteOperatorSetPrepare(r, Prepare);
  TfLiteOperatorSetInvoke(r, Eval);
  return r;
}

const TfLiteOperator* AtanOperator() {
  // Singleton instance, intentionally never destroyed.
  static const TfLiteOperator* atan_op = AtanOpCreate();
  return atan_op;
}

初始化 OpResolver 時，請將自訂作業新增至解析器 (請參閱下方的範例)。這會向 LiteRT 註冊運算子，以便 LiteRT 使用新的實作方式。

向核心程式庫註冊運算子

現在我們需要向核心程式庫註冊運算子。這項操作是透過 OpResolver 完成。在幕後，解譯器會載入核心程式庫，並指派該程式庫執行模型中的每個運算子。雖然預設程式庫只包含內建核心，但您可以替換/擴增程式庫，加入自訂程式庫 op 運算子。

OpResolver 類別會將運算子程式碼和名稱轉換為實際程式碼，定義方式如下：

class OpResolver {
 public:
  virtual TfLiteRegistration* FindOp(tflite::BuiltinOperator op) const = 0;
  virtual TfLiteRegistration* FindOp(const char* op) const = 0;
  ...
};

請注意，為了回溯相容性，這個類別使用較舊的具體型別 TfLiteRegistration，而非不透明型別 TfLiteOperator，但 TfLiteRegistration 結構體包含 TfLiteOperator* 型別的 registration_external 欄位。

MutableOpResolver 和 BuiltinOpResolver 類別衍生自 OpResolver：

class MutableOpResolver : public OpResolver {
 public:
  MutableOpResolver();  // Constructs an initially empty op resolver.
  void AddAll(const MutableOpResolver& other);
  ...
};

class BuiltinOpResolver : public MutableOpResolver {
 public:
  BuiltinOpResolver();  // Constructs an op resolver with all the builtin ops.
};

如要正常使用 (不含自訂作業)，您必須使用 BuiltinOpResolver 並編寫：

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

如要新增上述建立的自訂作業，可以改用 MutableOpResolver，並呼叫 tflite::AddOp (將解析器傳遞至 InterpreterBuilder 之前)：

tflite::ops::builtin::MutableOpResolver resolver;
resolver.AddAll(tflite::ops::builtin::BuiltinOpResolver());
tflite::AddOp(&resolver, AtanOpRegistration());

如果內建作業集過大，可以根據特定作業子集 (可能只有特定模型中包含的作業) 產生新的 OpResolver 程式碼。這相當於 TensorFlow 的選擇性註冊 (tools 目錄中提供簡單版本)。

如要在 Java 中定義自訂運算子，目前需要建構自己的自訂 JNI 層，並在 JNI 程式碼中編譯自己的 AAR。同樣地，如要在 Python 中定義這些運算子，可以將註冊項目放在 Python 包裝函式程式碼中。

請注意，如要支援一組作業而非單一運算子，可以按照上述類似程序操作。只要視需要新增任意數量的 AddCustom 運算子即可。此外，MutableOpResolver 也允許您使用 AddBuiltin 覆寫內建項目的實作。

測試及設定接線人員

如要使用 LiteRT 基準測試工具剖析作業，可以使用 LiteRT 的基準測試模型工具。為進行測試，您可以將適當的 AddCustom 呼叫 (如上所示) 新增至 register.cc，讓 LiteRT 的本機建構版本瞭解您的自訂作業。

最佳做法

  1. 請謹慎地最佳化記憶體配置和取消配置。在 Prepare 中分配記憶體比在 Invoke 中更有效率，而且在迴圈前分配記憶體比在每次疊代中分配記憶體更好。請使用暫時性張量資料，而非自行 malloc (請參閱項目 2)。盡可能使用指標/參照，而非複製。

  2. 如果資料結構會在整個作業期間持續存在，建議使用暫時性張量預先配置記憶體。您可能需要使用 OpData 結構體，參照其他函式中的張量索引。請參閱卷積的 Kernel 範例。程式碼片段範例如下。

    struct MyOpData {
  int temp_tensor_index;
  ...
};

void* Init(TfLiteOpaqueContext* context,
    const char* buffer, size_t length) {
  auto* op_data = new MyOpData{};
  ...
  return op_data;
}
void Free(TfLiteOpaqueContext* context, void* buffer) {
  ...
  delete reinterpret_cast<MyOpData*>(buffer);
}
TfLiteStatus Prepare(TfLiteOpaqueContext* context,
                     TfLiteOpaqueNode* node) {
  ...
  auto* op_data =
      reinterpret_cast<MyOpData*>(TfLiteOpaqueNodeGetUserData(node));
  const int num_temporaries = 1;
  int temporary_tensor_indices[num_temporaries];
  TfLiteOpaqueTensorBuilder* builder = TfLiteOpaqueTensorBuilderCreate();
  TfLiteOpaqueTensorBuilderSetType(builder, kTfLiteFloat32);
  TfLiteOpaqueTensorBuilderSetAllocationType(builder, kTfLiteArenaRw);
  TfLiteOpaqueContextAddTensor(context, builder,
      &temporary_tensor_indices[0]);
  TfLiteOpaqueTensorBuilderDelete(builder);
  TfLiteOpaqueNodeSetTemporaries(node, temporary_tensor_indices,
      num_temporaries);
  op_data->temp_tensor_index = temporary_tensor_indices[0];
  ...
  return kTfLiteOk;
}
TfLiteStatus Invoke(TfLiteOpaqueContext* context,
                    TfLiteOpaqueNode* node) {
  ...
  auto* op_data = reinterpret_cast<MyOpData*>(
      TfLiteOpaqueNodeGetUserData(node));
  TfLiteOpaqueTensor* temp_tensor =
      TfLiteOpaqueContextGetOpaqueTensor(context,
          op_data->temp_tensor_index);
  TF_LITE_OPAQUE_ENSURE(context,
      TfLiteTensorType(temp_tensor) == kTfLiteFloat32);
  TF_LITE_OPAQUE_ENSURE(context,
      TfLiteTensorGetAllocationType(temp_Tensor) == kTfLiteArenaRw);
  void *temp_data = TfLiteTensorData(temp_tensor);
  TF_LITE_OPAQUE_ENSURE(context, temp_data != nullptr);
  ...
  return kTfLiteOk;
}

  3. 如果不會浪費太多記憶體，建議使用靜態固定大小的陣列 (或 Resize 中的預先分配 std::vector)，而不是在每次執行疊代時動態分配 std::vector。

  4. 請避免例項化尚未存在的標準程式庫容器範本，因為這會影響二進位大小。舉例來說，如果作業需要使用其他核心沒有的 std::map，則可使用具有直接索引對應的 std::vector，同時維持較小的二進位大小。查看其他核心使用的內容，深入瞭解 (或提出問題)。

  5. 檢查 malloc 傳回的記憶體指標。如果這個指標是 nullptr，就不應使用該指標執行任何作業。如果您在函式中malloc，且有錯誤結束，請在結束前解除分配記憶體。

  6. 使用 TF_LITE_OPAQUE_ENSURE(context, condition) 檢查特定條件。使用 TF_LITE_OPAQUE_ENSURE 時，程式碼不得留下懸置記憶體，也就是說，應在分配任何會洩漏的資源前使用這些巨集。