# 自定义算子

TensorFlow Text 算子是自定义算子的示例。请参阅将 TF Text 转换为 TF Lite 教程了解代码示例。

## 示例：自定义 `Atan` 算子

### 创建 TensorFlow 模型

``````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
def train(x, y):
predicted_y = atan(x)
loss = tf.reduce_sum(tf.square(predicted_y - y))

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
``````

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

### 转换为 TensorFlow Lite 模型

```converter = tf.lite.TFLiteConverter.from_concrete_functions([atan.get_concrete_function()], atan)
<b>converter.allow_custom_ops = True</b>
tflite_model = converter.convert()
```

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

``````Encountered unresolved custom op: Atan.
``````

### 创建并注册算子

``````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;
``````

``````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
``````

### 在 TensorFlow Lite 运行时中定义内核

``````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;
}
``````

### 在内核库中注册算子

`OpResolver` 类会将算子代码和名称翻译成实际代码，其定义如下：

``````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;
};
``````

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

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

## 最佳做法

1. 谨慎优化内存分配和取消分配。在 `Prepare` 中分配内存比在 `Invoke` 中分配更高效，并且最好在循环之前而非在每次迭代中分配内存。使用临时张量数据，而不要自己分配内存（请参阅第 2 项）。使用指针/引用而不是无节制地进行复制。

2. 如果某个数据结构在整个运算期间持续存在，建议使用临时张量预分配内存。您可能需要使用 OpData 结构来引用其他函数中的张量索引。请参阅卷积内核中的示例。示例代码段如下：

``````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. 如果不想让它浪费太多内存，最好使用静态固定大小的数组（或在 `Resize` 中预分配的 `std::vector`），而不要使用在执行的每次迭代时动态分配的 `std::vector`

4. 避免实例化尚不存在的标准库容器模板，因为它们会影响二进制文件的大小。例如，如果您需要在运算中使用其他内核中不存在的 `std::map`，可以使用具有直接索引映射的 `std::vector`，同时保持较小的二进制文件大小。请查看其他内核使用的内容以获得深入见解（或询问）。

5. 检查指向由 `malloc` 返回的内存的指针。如果此指针是 `nullptr`，则不应使用该指针执行任何运算。如果在函数内 `malloc` 并出现退出错误，请在退出前释放内存。

6. 使用 `TF_LITE_ENSURE(context, condition)` 检查特定条件。使用 `TF_LITE_ENSURE` 时，您的代码不得将内存挂起（即，应该在分配任何可能泄漏的资源之前使用这些宏）。

[]
[]
{"lastModified": "\u6700\u540e\u66f4\u65b0\u65f6\u95f4 (UTC)\uff1a2024-01-11\u3002"}