自訂運算子

由於 LiteRT 內建運算子程式庫僅支援 TensorFlow 運算子的數量,不一定每個模型都可以轉換。詳情 請參閱運算子相容性

如要支援轉換,使用者可以自行提供 LiteRT 中不支援的 TensorFlow 運算子,也稱為自訂運算子。 如果您想合併多個不支援 (或支援) 的多個選項 將 TensorFlow 運算子轉換成單一融合式最佳化運算子,這是 運算子融合

自訂運算子包含四個步驟。

我們來看看透過自訂模型 執行模型的端對端範例 運算子 tf.atan (命名為 Atan,請參閱建立 TensorFlow 模型)。 TensorFlow 支援,但 LiteRT 不支援。

TensorFlow Text 運算子是自訂運算子的一個例子。詳情請參閱 將 TF 文字轉換為 LiteRT 教學課程,以取得程式碼範例。

範例:自訂 Atan 運算子

以下舉例說明支援 TensorFlow 運算子的 LiteRT 沒有。假設我們使用 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 來定義, 包含不透明類型 (TfLiteRegistrationExternal) 和相關函式。

TfLiteRegistrationExternal 是不透明類型:

typedef struct TfLiteRegistrationExternal TfLiteRegistrationExternal;

TfLiteRegistrationExternal 會儲存操作員的身分和實作項目。 (請注意,運算子與運算元不同,運算元儲存在 呼叫運算子的節點適用的 LiteRT 圖形節點)。

這種類型的實例是透過對 TfLiteRegistrationExternalCreate,可藉由呼叫 TfLiteRegistrationExternalDelete

運算子的身分是透過建構函式函式的參數來設定 TfLiteRegistrationExternalCreate:

TfLiteRegistrationExternal*
TfLiteRegistrationExternalCreate(
    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.
);

運算子實作可以定義「方法」並使用下列簽名。 這些方法都屬於選用項目,但如果運算子成功, 運算子實作方式必須 函式) 至少使用 PrepareInvoke 方法。

// 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);

運算實作中的函式名稱 (或 C++ 的命名空間前置字串) 不必與上述程式碼片段中的函式名稱相符,因為 TF Lite 自訂作業 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,因此這些「方法」在 TfLiteRegistrationExternal 類型,設定方法是透過傳遞 您的實作函式對應至對應的 setter 函式 TfLiteRegistrationExternalSetMethodName

void TfLiteRegistrationExternalSetInit(
    TfLiteRegistrationExternal* registration,
    void* (*init)(TfLiteOpaqueContext* context, const char* buffer,
                  size_t length));
void TfLiteRegistrationExternalSetFree(
    TfLiteRegistrationExternal* registration,
    void (*free)(TfLiteOpaqueContext* context, void* data));
void TfLiteRegistrationExternalSetPrepare(
    TfLiteRegistrationExternal* registration,
    TfLiteStatus (*prepare)(TfLiteOpaqueContext* context,
                            TfLiteOpaqueNode* node));
void TfLiteRegistrationExternalSetInvoke(
    TfLiteRegistrationExternal* registration,
    TfLiteStatus (*invoke)(TfLiteOpaqueContext* context,
                           TfLiteOpaqueNode* node));
void TfLiteRegistrationExternalSetAsyncKernel(
    TfLiteRegistrationExternal* registration,
    struct TfLiteAsyncKernel* (*async_kernel)(TfLiteOpaqueContext* context,
                                              TfLiteOpaqueNode* node));

詳情請參閱 common.h敬上 進一步瞭解 TfLiteContextTfLiteNode。「TfLiteContext」提供錯誤 報告設施及全域物件的存取權限,包括所有張量。 TfLiteNode 允許運算子實作存取其輸入和輸出內容。

解譯器載入模型時,會為每個模型呼叫一次 Init() 方法 節點如果運算是多次,系統就會多次呼叫特定 Init() 多次用於圖表中如果是自訂作業,則設定緩衝區會 ,內含可將參數名稱對應至其值的 Flexbuffer。 針對內建作業的緩衝區為空,因為解譯器已剖析 op 參數。需要狀態的核心實作項目應初始化 ,並將擁有權轉移給來電者。每呼叫 Init() 會有 對應的 Free() 呼叫,可讓實作 這些緩衝區可能在 Init() 中配置了

每當輸入張量調整大小時,直譯器就會經過 通知執行變更的圖表這樣他們才有機會 調整內部緩衝區的大小、檢查輸入形狀和類型是否有效,以及 重新計算輸出形狀這項操作是透過 Prepare() 方法完成 實作方式可以使用 TfLiteOpaqueNodeGetUserData(node)

最後,每次執行推論時,解譯器都會週遊圖形呼叫 Invoke() 方法,狀態也可用 TfLiteOpaqueNodeGetUserData(node)

定義這些「方法」即可實作自訂運算函式 定義會傳回 TfLiteRegistrationExternal 例項的函式 呼叫 TfLiteRegistrationExternalCreate,然後向相關函式 setter 方法:

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 TfLiteRegistrationExternal* MyCustomOpRegistrationExternal() {
    // Singleton instance, intentionally never destroyed.
    static const TfLiteRegistrationExternal* my_custom_op = ()[] {
        TfLiteRegistrationExternal* r =
            TfLiteRegistrationExternalCreate(
                kTfLiteBuiltinCustom, "MyCustomOp", /*version=*/ 1);
        TfLiteRegistrationExternalSetInit(r, Init);
        TfLiteRegistrationExternalSetFree(r, Free);
        TfLiteRegistrationExternalSetPrepare(r, Prepare);
        TfLiteRegistrationExternalSetInvoke(r, Eval);
        return r;
      };
    return my_custom_op;
  }

  const TfLiteRegistration* MyCustomOpRegistration() {
    static const TfLiteRegistration my_custom_op {
      .registration_external = MyCustomOpRegistrationExternal();
    };
    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 TfLiteRegistrationExternal* MyCustomOpCreate() {
  const TfLiteRegistrationExternal* r =
      TfLiteRegistrationExternalCreate(
          kTfLiteBuiltinCustom, "MyCustomOp", /*version=*/ 1);
  TfLiteRegistrationExternalSetInit(r, MyCustomOpInit);
  TfLiteRegistrationExternalSetFree(r, MyCustomOpFree);
  TfLiteRegistrationExternalSetPrepare(r, MyCustomOpPrepare);
  TfLiteRegistrationExternalSetInvoke(r, MyCustomOpEval);
  return r;
}

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

const TfLiteRegistration MyCustomOpRegistration() {
  static const TfLiteRegistration my_custom_op {
    .registration_external = MyCustomOpRegistrationExternal();
  };
  return my_custom_op;
}
      

請注意,註冊不會自動,並會對您的 應建立 MyCustomOpRegistration 函式 (詳情請參閱下方說明)。雖然 標準 BuiltinOpResolver (可從 :builtin_ops 目標取得) 就必須收集自訂作業 不同的自訂程式庫

在 LiteRT 執行階段中定義核心

如要在 LiteRT 中使用註解,我們要定義兩個函式 (PrepareEval) 和第三個呼叫建構 TfLiteRegistrationExternal

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 TfLiteRegistrationExternal* AtanOpRegistrationExternal() {
    // Singleton instance, intentionally never destroyed.
    static const TfLiteRegistrationExternal* atan_op = ()[] {
        auto* r = TfLiteRegistrationExternalCreate(
            kTfLiteBuiltinCustom, "ATAN", /*version=*/ 1);
        TfLiteRegistrationExternalSetPrepare(r, Prepare);
        TfLiteRegistrationExternalSetInvoke(r, Eval);
        return r;
      };
    return atan_op;
  }

  const TfLiteRegistration AtanOpRegistration() {
    static const TfLiteRegistration atan_op {
      .registration_external = AtanOpRegistrationExternal();
    };
    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 TfLiteRegistrationExternal* AtanOpCreate() {
  TfLiteRegistrationExternal* r = TfLiteRegistrationExternalCreate(
          kTfLiteBuiltinCustom, "ATAN", /*version=*/ 1);
  TfLiteRegistrationExternalSetPrepare(r, Prepare);
  TfLiteRegistrationExternalSetInvoke(r, Eval);
  return r;
}

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

const TfLiteRegistration AtanOpRegistration() {
  static const TfLiteRegistration atan_op {
    .registration_external = AtanOpRegistrationExternal();
  };
  return atan_op;
}
      

初始化 OpResolver 時,請將自訂運算新增至解析器中 (請參閱 )。這會向 LiteRT 註冊運算子 讓 LiteRT 可以使用新的實作方式請注意 TfLiteRegistration 中的引數對應 AtanPrepareAtanEval 您為自訂運算定義的函式。如果您先前使用 AtanInitAtanFree 函式,初始化運算中使用的變數並釋出空間。 那麼系統會將這些參數新增到 TfLiteRegistration;這些引數都設為 nullptr

使用核心程式庫註冊運算子

現在我們需要向核心程式庫註冊運算子。方法是使用 OpResolver。翻譯器會在背景載入 為執行模型中每個運算子而指派的核心。 雖然預設程式庫只包含內建核心,但也可以 將其替換為自訂程式庫運算運算子。

OpResolver 類別,會將運算子和名稱轉譯為實際 程式碼的定義如下:

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

請注意,為了兼顧回溯相容性,這個類別會使用舊版具體類型 TfLiteRegistration 而非不透明類型 TfLiteRegistrationExternal, 但 TfLiteRegistration 結構包含的 registration_external 欄位是 輸入 TfLiteRegistrationExternal*

MutableOpResolverBuiltinOpResolver 類別衍生自 OpResolver:

class MutableOpResolver : public OpResolver {
 public:
  MutableOpResolver();  // Constructs an initially empty op resolver.
  void AddBuiltin(tflite::BuiltinOperator op, const TfLiteRegistration* registration) = 0;
  void AddCustom(const char* op, const TfLiteRegistration* registration) = 0;
  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, 並呼叫 AddCustom (在將解析器傳遞至 InterpreterBuilder):

tflite::ops::builtin::MutableOpResolver resolver;
resolver.AddAll(tflite::ops::builtin::BuiltinOpResolver());
resolver.AddCustom("Atan", AtanOpRegistration());

如果系統認為這一組內建作業太大,系統可能會顯示新的 OpResolver 根據特定作業子集產生的程式碼,可能只有 特定模型的預測結果這相當於 TensorFlow 的選擇性註冊 (您可在 tools 目錄中找到簡易版本)。

如要用 Java 定義自訂運算子,目前您需要 建構自己的自訂 JNI 層,並編譯自己的 AAR 同樣地,如要以 Python 定義這些運算子, 以便將註冊資訊放在 Python 包裝函式程式碼

請注意,也可採用類似的程序來支援一組 而非單一運算子只要新增 AddCustom 運算子,數量不限 。此外,MutableOpResolver 還可讓您覆寫 使用 AddBuiltin 實作內建功能。

測試及剖析電信業者

如要透過 LiteRT 基準工具剖析運算,您可以使用 基準模型工具 LiteRT 也是個好方法如要進行測試,您可以將 新增適當的 AddCustom,以便 LiteRT 留意您的自訂運算 呼叫 (如上所示) register.cc

最佳做法

  1. 謹慎地最佳化記憶體配置及取消配置。正在分配記憶體 在 Prepare 中,比 Invoke 更有效率,並分配記憶體 會比每次疊代都更佳使用臨時張量資料 而不是賣掉自己的畫面 (見第 2 項)。改用指標/參照 盡可能減少複製的數量

  2. 如果資料結構在整個作業期間都會持續存在,建議您 使用臨時張量預先分配記憶體。您可能需要使用 OpData 結構體,用於參照其他函式中的張量索引。詳情請參閱 ,在 卷積核心。 下列是程式碼片段範例。

    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::mapstd::vector 搭配直接建立索引的對應關係,同時將 二進位檔的大小很小瞭解其他核心使用哪些功能取得深入分析資訊 (或提問)。

  5. 檢查 malloc 傳回記憶體的指標。如果此指標是 nullptr,不應使用該指標執行任何作業。如果發生以下情況: 函式中的 malloc 而且發生錯誤,因此在您之前,請先釋放記憶體 結束。

  6. 使用 TF_LITE_OPAQUE_ENSURE(context, condition) 檢查 值。您的程式碼不可因為 使用了 TF_LITE_OPAQUE_ENSURE,也就是這些巨集應 避免分配到流失的資源