4. 输出张量内存复用

任务目标

在RK3588 NPU连续推理场景中，每次调用rknn_run后获取输出张量时，若反复调用rknn_outputs_get重新分配内存，会产生大量堆内存碎片与分配开销。本案例实现输出张量内存复用：首次推理后保存输出缓冲区指针，后续推理直接复用该内存，仅更新数据内容。

环境依赖

• RK3588开发板，系统已安装librknnrt.so（RKNN Runtime）
• Python 3.8+，已安装rknn-toolkit2（用于模型转换，本案例仅需运行时API）
• C语言编译工具链（gcc / aarch64-linux-gnu-gcc）

核心实现

以下代码展示完整的输出张量内存复用流程。关键点：使用rknn_outputs_get获取输出缓冲区指针后，不调用rknn_outputs_release释放，而是保留该指针；下次推理时，同一输出张量索引对应的缓冲区会被Runtime自动覆盖写入新数据。通过rknn_query查询输出张量属性，确认缓冲区大小不变后即可安全复用。

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "rknn_api.h"

#define MODEL_PATH "./model.rknn"
#define INPUT_SIZE 224*224*3

static unsigned char *load_file(const char *path, int *size) {
    FILE *fp = fopen(path, "rb");
    if (!fp) return NULL;
    fseek(fp, 0, SEEK_END);
    *size = ftell(fp);
    fseek(fp, 0, SEEK_SET);
    unsigned char *buf = (unsigned char*)malloc(*size);
    fread(buf, 1, *size, fp);
    fclose(fp);
    return buf;
}

int main() {
    // 1. 初始化
    rknn_context ctx;
    int model_size = 0;
    unsigned char *model = load_file(MODEL_PATH, &model_size);
    if (rknn_init(&ctx, model, model_size, 0) < 0) {
        printf("rknn_init failed\n");
        return -1;
    }

    // 2. 查询输入/输出信息
    rknn_input_output_num io_num;
    rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
    printf("Input num: %d, Output num: %d\n", io_num.n_input, io_num.n_output);

    rknn_tensor_attr out_attr[io_num.n_output];
    memset(out_attr, 0, sizeof(out_attr));
    for (int i = 0; i < io_num.n_output; i++) {
        out_attr[i].index = i;
        rknn_query(ctx, RKNN_QUERY_OUT_ATTR, &out_attr[i], sizeof(out_attr[i]));
        printf("Output[%d] size: %d\n", i, out_attr[i].n_elems * 4); // 假设float32
    }

    // 3. 准备输入数据（示例：全零）
    rknn_input inputs[1];
    memset(inputs, 0, sizeof(inputs));
    inputs[0].index = 0;
    inputs[0].type = RKNN_TENSOR_UINT8;
    inputs[0].size = INPUT_SIZE;
    inputs[0].fmt = RKNN_TENSOR_NHWC;
    inputs[0].buf = calloc(1, INPUT_SIZE);

    // 4. 首次推理：获取输出缓冲区并保留
    rknn_run(ctx, NULL);
    rknn_output outputs[io_num.n_output];
    memset(outputs, 0, sizeof(outputs));
    for (int i = 0; i < io_num.n_output; i++) {
        outputs[i].index = i;
        outputs[i].want_float = 1;
    }
    rknn_outputs_get(ctx, io_num.n_output, outputs, NULL);
    printf("First inference: output buffers allocated at:\n");
    for (int i = 0; i < io_num.n_output; i++) {
        printf("  Output[%d] ptr = %p, size = %d\n", i, outputs[i].buf, outputs[i].size);
    }

    // 5. 连续推理10次，复用同一输出缓冲区
    for (int iter = 1; iter <= 10; iter++) {
        // 更新输入（示例：填充随机值）
        memset(inputs[0].buf, iter, INPUT_SIZE);
        rknn_inputs_set(ctx, 1, inputs);

        rknn_run(ctx, NULL);
        // 关键：不调用 rknn_outputs_release，直接复用 outputs[i].buf
        rknn_outputs_get(ctx, io_num.n_output, outputs, NULL);

        // 验证：检查指针是否与首次一致（即内存复用生效）
        if (iter == 1) {
            printf("Second inference ptr check:\n");
            for (int i = 0; i < io_num.n_output; i++) {
                printf("  Output[%d] ptr = %p (same as first: %s)\n",
                       i, outputs[i].buf,
                       (outputs[i].buf == outputs[i].buf) ? "yes" : "no");
            }
        }
    }

    // 6. 最终释放（仅一次）
    rknn_outputs_release(ctx, io_num.n_output, outputs);
    free(inputs[0].buf);
    free(model);
    rknn_deinit(ctx);
    printf("Memory reuse test completed.\n");
    return 0;
}

编译与运行

# 交叉编译（RK3588）
aarch64-linux-gnu-gcc -o output_reuse output_reuse.c -I./rknn_api -L./lib -lrknnrt

# 拷贝到开发板后运行
adb push output_reuse /userdata/
adb push model.rknn /userdata/
adb shell "cd /userdata && ./output_reuse"

验证方法

• 内存复用判定：程序打印的Output[i] ptr地址在10次推理中保持不变，说明未重新分配内存。
• 正确性检查：输出数据随输入变化（输入填充iter值），若输出张量数值同步变化，说明缓冲区内容被正确覆盖更新。
• 资源释放：程序结束时仅调用一次rknn_outputs_release，无内存泄漏（可用valgrind或长期运行观察内存稳定）。

注意事项

• 复用输出缓冲区的前提是模型输出张量形状固定。若模型支持动态shape，需在每次推理前查询RKNN_QUERY_OUT_ATTR确认n_elems未变化，否则需重新分配。
• 本案例未调用rknn_outputs_release，因此Runtime内部不会释放输出缓冲区。若需中途释放，可调用后重新rknn_outputs_get获取新缓冲区。
• 多线程场景下，每个线程需独立持有输出缓冲区指针，不可共享复用。