模型部署——使用rknn-toolkit将Pytorch模型转为RKNN模型（详细图文教程）

Source

在这里插入图片描述

一、环境准备

在进行RKNN模型转换之前，需要安装rknn-toolkit工具，关于此工具的具体安装过程参考我另外一篇博文：安装rknn-toolkit工具

下面是我自己运行的环境，供参考：
在这里插入图片描述

二、源码包获取

我提供一个resnet18的模型和测试数据集，还有相关转换代码，都打包在源码包中，，获取方法为文章末扫码到公众号中回复关键字：模型部署转RKNN模型。获取下载链接。

下载解压后的样子如下：

在这里插入图片描述

三、RKNN模型转换

3.1 模型转换的流程图

在这里插入图片描述

3.2 参数修改

在这里插入图片描述

在我提供的源码包中，直接运行脚本export_rknn.py即可。

3.3 代码

转换代码为：

## 测试用来构建RKNN模型的API

from rknn.api import RKNN

if __name__=="__main__":
    rknn = RKNN(verbose=True,verbose_file="log.txt")   # verbose为True表示打印详细的日志，verbose_file表示将日志存放到指定的路径中

    # 调用config接口配置要生成的RKNN模型
    # 调用config接口设置模型的预处理、量化方法等参数
    rknn.config(
        mean_values = [[123.675,116.28,103.53]],        # mean_values表示预处理要减去的均值化参数
        std_values = [[58.395,58.395,58.395]],          # std_values 表示预处理要除的标准化参数
        quantized_dtype = "asymmetric_quantized-8",     # 表示默认为8位非对称量化。quantized_dtype表示量化类型   通用的模型权重和激活值都是float32类型的，会占据4个字节，而经过8位的非对称量化之后，权重和激活值量化为int8类型，只占1个字节
        quantized_algorithm = "normal",                   # quantized_algorithm表示量化的算法，目前支持norm,mmse,kl
        quantized_method = "channel",                   # quantized_method表示量化方式,总的有channel（通道级量化）和layer（层级量化）两种，channel的精度要要高一些，默认为channel
        quant_img_RGB2BGR = False,                      # 表示在做量化时，是否做RGB2BGR的转换。此参数只会应用到量化阶段，并不会嵌入到RKNN模型中。
        target_platform = "rk3588",                     # target_platform表示生成的RKNN模型要运行在哪个RKNPU平台上。通常有rk3588,rk3566,rv1126等
        # target_platform="rv1126",
        float_dtype = "float16",                        # 表示RKNN模型的浮点数的类型，目前只支持float16的格式。如果不进行量化操作，则会将原始的float32格式默认转为float16格式。
        optimization_level = 3,                         # 表示优化等级，默认为3级，表示打开全部优化选项。如果设置为0则表示关闭所有优化选项。1，2代表中间值，表示打开部分优化选项。会对最后的精度值产生影响
        custom_string = "this is my rknn model!",       # 表示向RKNN模型中添加的自定义字符串信息，可以在CAPI中通过查询接口，查询到添加的自定义字符串信息
        remove_weight = False,                          # 表示生成一个去除权重的从模型，可以使用CAPI与另一个完整的模型共享权重，从而减小内存的消耗 一般用在rv109和rv106上
        compress_weight = False,                        # 压缩权重，可以减少模型的体积
        inputs_yuv_fmt = None,                          # 表示RKNN模型输入数据的yuv格式
        single_core_mode = False,                       # 表示构建的RKNN模型运行在核心模式，只适用于RK3588
    )

    # 添加load_xxx接口，进行常用深度学习模型的导入           将深度学习模型导入
    rknn.load_pytorch(
        model = "resnet18.pt",            # model表示加载模型的地址
        input_size_list = [[1,3,224,224]],                            # 表示模型输入节点对应图片的尺寸和通道数
    )

    # 使用build接口来构建RKNN模型
    rknn.build(
        do_quantization=True,                      # 表示是否对生成的rknn模型进行量化操作
        dataset = "dataset.txt",                   # 此参数只有上面的do_quantization为True时才生效，表示RKNN模型量化时采用的量化校正数据集   dataset.txt中填写的就是要量化的图片（正常情况下有20-100张量化图片）
        rknn_batch_size = -1,                      # 调整输入的batch数量，可以提高推理的效率
    )

    # 调用export_rknn接口导出RKNN模型
    rknn.export_rknn(
        export_path="resnet18.rknn"              # 表示导出的RKNN模型路径和名称

    )

    rknn.release()  # 释放

3.4 运行结果

运行脚步后的输出结果如下：
在这里插入图片描述

3.5 转换得到的RKNN模型

转换后得到的RKNN模型如下：

在这里插入图片描述

四、总结

以上就是使用rknn-toolkit将Pytorch模型转为RKNN模型的详细过程，希望能帮到你！

总结不易，多多支持，谢谢！

目录

一、环境准备

二、源码包获取

三、RKNN模型转换

3.1 模型转换的流程图

3.2 参数修改

3.3 代码

3.4 运行结果

3.5 转换得到的RKNN模型

四、总结