自定义运算符

由于 LiteRT 内置运营商库仅支持有限的 TensorFlow 运算符数量，但并非所有模型都是可转换的。如需了解详情，请参阅 请参阅运算符兼容性。

要允许转换，用户可以提供自定义的 不支持的 TensorFlow 运算符，称为自定义运算符。 如果您希望将一系列不支持的（或受支持的） 将 TensorFlow 运算符转换为单个经过优化的自定义运算符，请参阅 运算符融合。

使用自定义运算符包含四个步骤。

我们来看一个端到端示例，了解如何使用自定义 运算符 tf.atan（名为 Atan，请参阅创建 TensorFlow 模型），其中 在 TensorFlow 中受支持，但在 LiteRT 中不受支持。

TensorFlow Text 运算符就是一个自定义运算符的一个示例。请参阅 如需查看代码示例，请参阅将 TF Text 转换为 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.
);

运算符实现可以定义“方法”签名。 所有这些方法都是可选的，但要让操作员成功执行这些操作， 运算符实现需要定义和设置（使用 setter 函数）至少包含 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);

操作实现中的函数 names（对于 C++，则为命名空间前缀） 不必与上述代码段中的函数名称一致，因为 TF 精简版自定义操作 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，因此这些“方法”以 C 函数指针的形式实现 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 ，详细了解TfLiteContext和TfLiteNode。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 中使用相应操作，我们只需定义两个函数 （Prepare 和 Eval），第三个函数用于构造 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 中的参数对应于 AtanPrepare 和 AtanEval 为自定义操作定义的函数。如果您使用的是 AtanInit 和 AtanFree 函数来初始化操作中使用的变量和释放空间， 那么它们会分别添加到 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*。

MutableOpResolver 和 BuiltinOpResolver 类派生自 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 基准工具分析操作，您可以使用 基准模型工具 。出于测试目的，您可以将 添加适当的 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::map 具有直接索引映射的 std::vector 可以正常运行，同时 因为二进制文件很小了解其他内核使用什么来获得数据洞见（或提问）。

  5. 检查指向由 malloc 返回的内存的指针。如果此指针的状态是 nullptr，则不应使用该指针执行任何操作。如果您 malloc 并且退出了错误，请先取消分配内存， 退出。

  6. 使用 TF_LITE_OPAQUE_ENSURE(context, condition) 检查特定 条件。在以下情况下，代码不得使内存挂起 TF_LITE_OPAQUE_ENSURE，也就是说，应该先使用这些宏 分配的所有会泄露的资源