📚 课程目录第2章内存分配策略初步优化

2. 内存分配策略初步优化

任务目标

在RK3588 NPU推理场景中，频繁调用rknn_init/rknn_run时默认内存分配策略会产生碎片，导致峰值内存上升。本案例通过设置rknn_set_mem_strategy为顺序分配模式，并配合输入/输出张量的显式预分配，减少碎片、提升利用率。

环境依赖

RK3588开发板，系统为Ubuntu 20.04+
RKNN Toolkit 1.6.0（PC端模型转换）
RKNN Runtime 1.6.0（板端推理库，含librknnmrt.so）
模型文件：mobilenet_v2.rknn（由RKNN Toolkit转换得到）

关键实现

核心策略：在rknn_init后立即调用rknn_set_mem_strategy，将分配策略从默认的RKNN_MEM_STRATEGY_DEFAULT切换为RKNN_MEM_STRATEGY_SEQUENTIAL。该模式按推理顺序连续分配内存块，避免反复malloc/free造成的空洞。同时，对输入/输出张量使用rknn_create_mem预分配，使NPU内部与外部张量共享同一内存池，减少拷贝与碎片。

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

#define MODEL_PATH "./mobilenet_v2.rknn"
#define IMG_SIZE 3 * 224 * 224

int main() {
    // 1. 加载模型
    FILE *fp = fopen(MODEL_PATH, "rb");
    if (!fp) { printf("open model failed\n"); return -1; }
    fseek(fp, 0, SEEK_END);
    int model_len = ftell(fp);
    fseek(fp, 0, SEEK_SET);
    void *model_data = malloc(model_len);
    fread(model_data, 1, model_len, fp);
    fclose(fp);

    // 2. 初始化NPU上下文
    rknn_context ctx;
    int ret = rknn_init(&ctx, model_data, model_len, 0, NULL);
    if (ret < 0) { printf("rknn_init failed: %d\n", ret); return -1; }

    // 3. 设置内存分配策略为顺序模式（关键步骤）
    ret = rknn_set_mem_strategy(ctx, RKNN_MEM_STRATEGY_SEQUENTIAL);
    if (ret < 0) {
        printf("set mem strategy failed: %d\n", ret);
        rknn_destroy(ctx);
        return -1;
    }

    // 4. 查询输入/输出张量信息
    rknn_input_output_num io_num;
    rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
    rknn_tensor_attr input_attr, output_attr;
    memset(&input_attr, 0, sizeof(input_attr));
    input_attr.index = 0;
    rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &input_attr, sizeof(input_attr));
    memset(&output_attr, 0, sizeof(output_attr));
    output_attr.index = 0;
    rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &output_attr, sizeof(output_attr));

    // 5. 预分配输入/输出内存（避免推理时临时分配）
    rknn_tensor_mem *input_mem = rknn_create_mem(ctx, input_attr.size);
    rknn_tensor_mem *output_mem = rknn_create_mem(ctx, output_attr.size);
    if (!input_mem || !output_mem) {
        printf("create mem failed\n");
        rknn_destroy(ctx);
        return -1;
    }

    // 6. 模拟连续推理10次，观察内存稳定性
    rknn_input inputs[1];
    rknn_output outputs[1];
    inputs[0].index = 0;
    inputs[0].type = RKNN_TENSOR_UINT8;
    inputs[0].size = IMG_SIZE;
    inputs[0].fmt = RKNN_TENSOR_NHWC;
    inputs[0].buf = malloc(IMG_SIZE);
    memset(inputs[0].buf, 0, IMG_SIZE);  // 用零填充模拟输入

    outputs[0].index = 0;
    outputs[0].want_float = 1;

    for (int i = 0; i < 10; i++) {
        // 使用预分配内存作为输入
        memcpy(input_mem->virt_addr, inputs[0].buf, IMG_SIZE);
        rknn_set_io_mem(ctx, input_mem, &input_attr);
        rknn_set_io_mem(ctx, output_mem, &output_attr);

        ret = rknn_run(ctx, NULL);
        if (ret < 0) { printf("rknn_run failed at iter %d: %d\n", i, ret); break; }

        ret = rknn_outputs_get(ctx, 1, outputs, NULL);
        if (ret < 0) { printf("get output failed: %d\n", ret); break; }
        rknn_outputs_release(ctx, 1, outputs);
    }

    // 7. 清理
    free(inputs[0].buf);
    rknn_free_mem(ctx, input_mem);
    rknn_free_mem(ctx, output_mem);
    rknn_destroy(ctx);
    free(model_data);
    printf("inference loop completed\n");
    return 0;
}

编译与运行

# 板端交叉编译（RK3588）
aarch64-linux-gnu-gcc mem_strategy_demo.c -o mem_strategy_demo \
    -I/usr/include/rknn -L/usr/lib -lrknnmrt

# 将可执行文件与模型拷贝至板端后运行
./mem_strategy_demo

验证方法

在板端运行程序前，先通过cat /proc/meminfo记录可用内存基线。程序运行期间，在另一个终端循环执行cat /proc/meminfo | grep MemFree观察内存波动。判定通过标准：

程序输出inference loop completed，无错误码；
10次推理过程中，MemFree值下降幅度不超过20MB（对比未设置策略时通常下降50MB以上）；
连续运行3次程序，内存峰值基本一致，无递增趋势。

若需对比优化效果，可将第3步的rknn_set_mem_strategy调用注释掉重新编译运行，观察内存波动差异。

常见问题

策略设置失败（返回非0）：确认RKNN Runtime版本≥1.6.0，旧版本不支持该接口。
内存预分配失败：检查input_attr.size是否与实际输入尺寸匹配，模型输入为224x224x3时该值应为150528。
性能下降：顺序分配策略在极端动态shape场景下可能增加单次分配耗时，但本案例固定shape场景下影响可忽略。