一、为什么需要自定义操作？

NNAPI 虽然内置了不少操作，但说实话，它不可能覆盖所有场景。我在做 AI 模型落地时，经常遇到一些「奇奇怪怪」的算子——比如某个模型里有个自定义的激活函数，或者一个特殊的池化方式。这时候怎么办？

嗯，NNAPI 给了我们一条路：自定义操作（Custom Operation）。说白了，就是你自己写一个操作，然后告诉 NNAPI：「嘿，这个操作我来处理，你别管了。」

我个人习惯把自定义操作分成两类：

• 纯计算型：比如自定义激活函数、自定义归一化
• 硬件加速型：比如调用 DSP、NPU 的私有指令

今天我们就重点讲第一类——如何编写、注册一个自定义操作，以及它的使用限制。

二、自定义操作的整体流程

先看一张图，帮你快速建立整体认知。我画了一个流程图，展示从定义到调用的完整链路：

这个流程看起来有点长，但别怕。我拆开来讲，你跟着做一遍就明白了。

三、第一步：定义操作类型

首先，你得告诉 NNAPI：「我要定义一个操作，它的编号是多少，输入输出是什么。」

NNAPI 预留了一个范围给自定义操作：ANEURALNETWORKS_OPERATION_CUSTOM_START 到 ANEURALNETWORKS_OPERATION_CUSTOM_END。你在这个范围内选一个值就行。

// 自定义操作类型定义
enum {
    MY_CUSTOM_OP_RELU_CUSTOM = ANEURALNETWORKS_OPERATION_CUSTOM_START,
    MY_CUSTOM_OP_TANH_CUSTOM,
    // 可以继续加...
};

我个人习惯把自定义操作的枚举值从 ANEURALNETWORKS_OPERATION_CUSTOM_START 开始递增，这样不容易冲突。你想想看，如果随便写个 0x1000，万一哪天 NNAPI 扩展了内置操作，撞上了就麻烦了。

四、第二步：实现回调函数

这是最核心的部分。NNAPI 在运行时遇到你的自定义操作，会调用你注册的回调。回调的签名是固定的：

typedef int (*ANeuralNetworksExecutionCallback)(
    ANeuralNetworksExecution *execution,
    void *userData,
    const ANeuralNetworksOperandType *inputTypes,
    uint32_t inputCount,
    ANeuralNetworksOperandType **outputTypes,
    uint32_t *outputCount
);

嗯，参数有点多。我解释一下：

• execution：当前执行的上下文，你可以用它来读写数据
• userData：你注册时传的自定义数据，比如一些配置参数
• inputTypes：输入张量的类型信息
• outputTypes：输出张量的类型信息（你需要填充它）

举个例子，假设我们要实现一个自定义的 ReLU 变体——不是简单的 max(0, x)，而是 max(alpha, x)，其中 alpha 是一个可配置的参数。

// 自定义 ReLU 的实现
int myReluCustomCallback(
    ANeuralNetworksExecution *execution,
    void *userData,
    const ANeuralNetworksOperandType *inputTypes,
    uint32_t inputCount,
    ANeuralNetworksOperandType **outputTypes,
    uint32_t *outputCount
) {
    // 1. 检查输入
    if (inputCount != 1) return ANEURALNETWORKS_BAD_DATA;
    
    // 2. 获取输入数据
    void *inputData = nullptr;
    uint32_t inputLength = 0;
    ANeuralNetworksExecution_getInputOperandData(execution, 0, &inputData, &inputLength);
    
    // 3. 获取用户参数（alpha 值）
    float alpha = userData ? *(float*)userData : 0.0f;
    
    // 4. 计算输出
    float *input = (float*)inputData;
    uint32_t elementCount = inputLength / sizeof(float);
    
    // 分配输出内存
    float *output = (float*)malloc(inputLength);
    for (uint32_t i = 0; i < elementCount; i++) {
        output[i] = input[i] > alpha ? input[i] : alpha;
    }
    
    // 5. 设置输出
    ANeuralNetworksExecution_setOutputOperandData(execution, 0, output, inputLength);
    
    free(output);
    return ANEURALNETWORKS_NO_ERROR;
}

我在项目中遇到过一个问题：回调里直接 malloc 分配内存，如果频繁调用，内存碎片会很严重。后来我改用了一个预分配的内存池，性能提升了不少。这个坑你可以提前避开。

五、第三步：注册到 NNAPI 模型

回调写好了，怎么告诉 NNAPI 呢？通过 ANeuralNetworksModel_addOperation 的扩展参数。

// 注册自定义操作
ANeuralNetworksModel *model;
// ... 创建 model 的代码 ...

// 定义操作输入
uint32_t inputIndex = 0;
ANeuralNetworksOperandType inputType = {
    .type = ANEURALNETWORKS_TENSOR_FLOAT32,
    .dimensionCount = 1,
    .dimensions = &inputDim,
    .scale = 0.0f,
    .zeroPoint = 0
};
ANeuralNetworksModel_addOperand(model, &inputType);

// 定义操作输出
uint32_t outputIndex = 1;
ANeuralNetworksOperandType outputType = {
    .type = ANEURALNETWORKS_TENSOR_FLOAT32,
    .dimensionCount = 1,
    .dimensions = &outputDim,
    .scale = 0.0f,
    .zeroPoint = 0
};
ANeuralNetworksModel_addOperand(model, &outputType);

// 添加自定义操作
uint32_t inputs[] = {inputIndex};
uint32_t outputs[] = {outputIndex};

// 关键：注册回调
float alpha = 0.1f;
ANeuralNetworksModel_addOperation(
    model,
    MY_CUSTOM_OP_RELU_CUSTOM,  // 自定义操作类型
    inputs, 1,
    outputs, 1,
    myReluCustomCallback,      // 回调函数
    &alpha,                    // userData
    sizeof(float)              // userData 大小
);

注意看最后三个参数：回调函数、用户数据、用户数据大小。这就是 NNAPI 在运行时调用你的入口。

六、使用限制与避坑指南

自定义操作虽然灵活，但限制也不少。我总结了几条：

七、一个完整的例子

最后，我放一个完整的流程代码片段，把上面所有步骤串起来：

// 1. 定义操作类型
#define MY_OP_CUSTOM_SIGMOID (ANEURALNETWORKS_OPERATION_CUSTOM_START + 0)

// 2. 实现回调
int mySigmoidCallback(ANeuralNetworksExecution *execution, void *userData,
                      const ANeuralNetworksOperandType *inputTypes, uint32_t inputCount,
                      ANeuralNetworksOperandType **outputTypes, uint32_t *outputCount) {
    void *inData; uint32_t inLen;
    ANeuralNetworksExecution_getInputOperandData(execution, 0, &inData, &inLen);
    
    float *in = (float*)inData;
    uint32_t n = inLen / sizeof(float);
    float *out = (float*)malloc(inLen);
    
    for (uint32_t i = 0; i < n; i++) {
        out[i] = 1.0f / (1.0f + expf(-in[i]));
    }
    
    ANeuralNetworksExecution_setOutputOperandData(execution, 0, out, inLen);
    free(out);
    return ANEURALNETWORKS_NO_ERROR;
}

// 3. 注册到模型
ANeuralNetworksModel_addOperation(model, MY_OP_CUSTOM_SIGMOID,
                                  inputs, 1, outputs, 1,
                                  mySigmoidCallback, NULL, 0);

// 4. 编译执行（和普通模型一样）
ANeuralNetworksCompilation *compilation;
ANeuralNetworksCompilation_create(model, &compilation);
ANeuralNetworksCompilation_finish(compilation);

ANeuralNetworksExecution *execution;
ANeuralNetworksExecution_create(compilation, &execution);
// 设置输入、执行、获取输出...

你看，其实没那么复杂。核心就是三步：定义类型、实现回调、注册。剩下的编译执行流程，和普通 NNAPI 模型完全一样。

嗯，关于自定义操作，今天就讲到这里。记住一点：自定义操作是最后的手段，能用内置操作组合解决的，就别自己造轮子。但如果真的需要，NNAPI 给了你足够的灵活性。