#AI挑战营第二站# Ubuntu下训练pytorch模型导出ONNX及RKNN转化

tinnu

#AI挑战营第二站# Ubuntu下训练pytorch模型导出ONNX及RKNN转化 [复制链接]

<div class='showpostmsg'><i class="pstatus"> 本帖最后由 tinnu 于 2024-5-27 19:21 编辑 </i>

## rknn 环境搭建
> 由于我使用的是markdown编辑器，发表帖子后缩进消失，代码请直接查看附件或仓库

### conda安装
- 我所使用的ubuntu kylin 20.04.3 LTS 实体机系统。python版本为3.8,虽然也使用conda托管，但还是沿用3.8.10经典版本。
```shell
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
bash Miniconda3-latest-Linux-x86_64.sh
```

### 创建conda环境
```shell
export PATH="/home/anaconda3/bin":$PATH
source /home/anaconda3/bin/activate
conda create -n torch1.10_rknn2 python=3.8
source /home/anaconda3/bin/activate torch1.10_rknn2
conda activate torch1.10_rknn2
```

### introduct
- Rockchip 的 RKNN(Rockchip Neural Network) 是一款 AI 推理框架，能够在 Rockchip 的 SoC 上调用芯片内部的npu加速单元。它提供了两个工具：RKNN-Toolkits 和 RKNPU。

- RKNPU 工具提供了运行时和一些demo，其中包含一些已经转化为.rknn的模型；在rknn2组件中，运行时是一个rknn_server程序，而rknn1中不需要该运行时支持，仅提供了驱动npu的ko文件，可以用于更新npu的驱动。不过出厂镜像默认支持，这里就不作更新。

- 另一个工具是 RKNN-Toolkits，这是一组转换模型工具，可以用于将常见的深度学习模型转换为 RKNN 支持的模型格式。RKNN-Toolkits 提供了一些脚本和 API，可以方便地将 Caffe、TensorFlow 等深度学习框架的模型转换为 RKNN 支持的格式，并将其部署到 RKNN 平台上进行推理运算。

- RKNN又分为RKNN1和RKNN2，RKNN1 是基于 NPUv1 架构实现的，而 RKNN2 是基于新一代 NPUv2 架构实现的。其两者SDK不兼容。我们要使用的RV1106的RKNPU是基于RKNN2，而用于对比的RV1126则是基于RKNN1。

### 安装 rknn_toolkit2
- 移除其他版本
    - pip3 uninstall rknn_toolkit
    - pip3 uninstall rknn_toolkit2
- 最新版本的 rknn_toolkit2 已经是1.6.0版本。这个版本在 rknn_toolkit2 仓库里面集成了 rknn_toolkit2 rknn_toolkit2_lite 和 rknpu，只需要下着一个即可。

```shell
    git clone https://github.com/rockchip-linux/rknn-toolkit2.git --depth=1
    pip3 install -r rknn-toolkit2/packages/requirements_cp38-1.6.0.txt
    pip3 install rknn-toolkit2/packages/rknn_toolkit2-1.6.0+81f21f4d-cp38-cp38-linux_x86_64.whl
    pip3 list | grep rknn
    ```

<ignore_js_op>

<div class="tip tip_4 aimg_tip" id="aimg_811856_menu" style="position: absolute; display: none" disautofocus="true">
<div class="xs0">
<p><strong>rknn安装.png</strong> <em class="xg1">(35.29 KB, 下载次数: 0)</em></p>
<p>
<a href="javascript:;" class="loginf" >下载附件</a>

<a href="javascript:;" onclick="showWindow(this.id, this.getAttribute('url'), 'get', 0);" id="savephoto_811856" url="home.php?mod=spacecp&ac=album&op=saveforumphoto&aid=811856&handlekey=savephoto_811856">保存到相册</a>

</p>

</div>
<div class="tip_horn"></div>
</div>

</ignore_js_op>

## 转化模型
- 之前帖子演示了 pytorch->onnx 模型转化，但这个转出来的模型进行 onnx->rknn ，会报错
    ```shell
    The input shape ['batch_size', 28, 28] of 'input' is not support!
          Please set the 'inputs' / 'input_size_list' parameters of 'rknn.load_onnx', or set the 'dyanmic_input' parameter of 'rknn.config' to fix the input shape!
    ```
    ```py
    input_example = torch.randn(1, 28, 28)
    model.eval()
    torch.onnx.export(
            model,     # The model to be exported
            input_example, # Example input data
            "model/model-3.onnx",   # Output file path
            opset_version=19,# ONNX operator set version (choose a suitable version based on your PyTorch version and model requirements)
            do_constant_folding=True,# Whether to execute constant folding for optimization
            input_names=["input"], # The model's input names
            output_names=["output"],# The model's output names
            dynamic_axes={
                "input": {0: "batch_size"},# Make batch size dynamic (allow different batch sizes at inference time)
                "output": {0: "batch_size"}# Similarly, make output batch size dynamic
            },
        )
    ```
- 经过一番处理，终于发现是 pytorch->onnx 转化时 dyanmic_input 的锅，删掉即可：
    ```py
    input_example = torch.randn(1, 1, 28, 28)
    model.eval()
    torch.onnx.export(
            model,     # The model to be exported
            input_example, # Example input data
            "../3.tran/model8.onnx",   # Output file path
            opset_version=11,# ONNX operator set version (choose a suitable version based on your PyTorch version and model requirements)
            do_constant_folding=True,# Whether to execute constant folding for optimization
        )
    ```
- onnx->rknn 转化脚本
    ```py
    # %%
    from rknn.api import RKNN

# 创建RKNN对象
    rknn = RKNN()
    rknn.config(mean_values=[[128]], std_values=[[128]], target_platform='rv1106')

# %% Load model
    import os

ONNX_MODEL = 'model8.onnx'
    RKNN_MODEL = 'model8.rknn'

if not os.path.exists(ONNX_MODEL):
        print('NO ONNX model...')
    else:
        print('--> Loading model')
        ret = rknn.load_onnx(model=ONNX_MODEL)
        if ret != 0:
            print('Load model failed!')
            exit(ret)
        print('done')

# %% Build model
    print('Building RKNN model...')
    rknn.build(do_quantization=True, dataset='subset_data.csv')

# %% Export rknn model
    print('--> Export rknn model')
    ret = rknn.export_rknn(export_path=RKNN_MODEL, cpp_gen_cfg=True)
    if ret != 0:
        print('Export rknn model failed!')
        exit(ret)
    print('done')
    ```
- 但这里有个问题，由于只有3588才支持不量化（do_quantization=True），必须量化的话，需要 dataset 文件，dataset生成需要两步，我把 datasetfile_com.py 脚本放进附录和工程里面。
    1. 转化数据集为图片文件
    2. 生成 dataset.csv
- 转化结束会报一些错误，但能征程生成rknn模型，暂时不管。

<ignore_js_op>

<div class="tip tip_4 aimg_tip" id="aimg_811855_menu" style="position: absolute; display: none" disautofocus="true">
<div class="xs0">
<p><strong>rknn转化.png</strong> <em class="xg1">(107.44 KB, 下载次数: 0)</em></p>
<p>
<a href="javascript:;" class="loginf" >下载附件</a>

<a href="javascript:;" onclick="showWindow(this.id, this.getAttribute('url'), 'get', 0);" id="savephoto_811855" url="home.php?mod=spacecp&ac=album&op=saveforumphoto&aid=811855&handlekey=savephoto_811855">保存到相册</a>

</p>

</div>
<div class="tip_horn"></div>
</div>

</ignore_js_op>

- 转化后还可以评估模型，这里不赘述
- [mpu_rv1106_rknn](https://gitlink.org.cn/tinnu/mpu_rv1106_rknn) 工程里面model文件夹有所有的转化脚本文件：

<ignore_js_op>

<a href="javascript:;" class="loginf">train_3.py</a>

<em class="xg1">(3.72 KB, 下载次数: 0)</em>
</span>
<div class="tip tip_4" id="attach_811835_menu" style="position: absolute; display: none" disautofocus="true">
<div class="tip_c xs0">
<div class="y">2024-5-27 18:01 上传</div>
点击文件名下载附件
训练模型，生成onnx
</div>
<div class="tip_horn"></div>
</div>

</ignore_js_op>

<ignore_js_op>

<a href="javascript:;" class="loginf">translate_1.py</a>

<em class="xg1">(768 Bytes, 下载次数: 0)</em>
</span>
<div class="tip tip_4" id="attach_811836_menu" style="position: absolute; display: none" disautofocus="true">
<div class="tip_c xs0">
<div class="y">2024-5-27 18:01 上传</div>
点击文件名下载附件
将onnx转化为rknn
</div>
<div class="tip_horn"></div>
</div>

</ignore_js_op>

<ignore_js_op>

<em class="xg1">(2.9 KB, 下载次数: 1)</em>
</span>
<div class="tip tip_4" id="attach_811837_menu" style="position: absolute; display: none" disautofocus="true">
<div class="tip_c xs0">
<div class="y">2024-5-27 18:02 上传</div>
点击文件名下载附件
转化出 dataset 文件
</div>
<div class="tip_horn"></div>
</div>

</ignore_js_op>

<ignore_js_op>

<a href="javascript:;" class="loginf">subset_data.csv</a>

<em class="xg1">(709 Bytes, 下载次数: 0)</em>
</span>
<div class="tip tip_4" id="attach_811834_menu" style="position: absolute; display: none" disautofocus="true">
<div class="tip_c xs0">
<div class="y">2024-5-27 18:01 上传</div>
点击文件名下载附件
dataset 文件
</div>
<div class="tip_horn"></div>
</div>

</ignore_js_op>

<ignore_js_op>

<a href="javascript:;" class="loginf">model8.rknn</a>

<em class="xg1">(50.73 KB, 下载次数: 0)</em>
</span>
<div class="tip tip_4" id="attach_811833_menu" style="position: absolute; display: none" disautofocus="true">
<div class="tip_c xs0">
<div class="y">2024-5-27 18:01 上传</div>
点击文件名下载附件
最终生成的rknn模型
</div>
<div class="tip_horn"></div>
</div>

</ignore_js_op>

<ignore_js_op>

<a href="javascript:;" class="loginf">model8.onnx</a>

<em class="xg1">(86.47 KB, 下载次数: 0)</em>
</span>
<div class="tip tip_4" id="attach_811832_menu" style="position: absolute; display: none" disautofocus="true">
<div class="tip_c xs0">
<div class="y">2024-5-27 18:01 上传</div>
点击文件名下载附件
训练生成的onnx模型
</div>
<div class="tip_horn"></div>
</div>

</ignore_js_op>

### 问题处理：
- 进行 onnx->rknn 模型转化，会报错：
    ```shell
    The input shape ['batch_size', 28, 28] of 'input' is not support!
          Please set the 'inputs' / 'input_size_list' parameters of 'rknn.load_onnx', or set the 'dyanmic_input' parameter of 'rknn.config' to fix the input shape!
    ```
- 这是由于 pytorch->onnx 转化时，设置了 dyanmic_input 的问题，去掉 dyanmic_input 即可成功

</div><script>					var loginstr = '<div class="locked">查看本帖全部内容，请<a href="javascript:;"   style="color:#e60000" class="loginf">登录</a>或者<a href="https://bbs.eeworld.com.cn/member.php?mod=register_eeworld.php&action=wechat" style="color:#e60000" target="_blank">注册</a></div>';
					
					if(parseInt(discuz_uid)==0){
						
					}		</script><script type="text/javascript">(function(d,c){var a=d.createElement("script"),m=d.getElementsByTagName("script")[0],eewurl="//counter.eeworld.com.cn/pv/count/";a.src=eewurl+c;m.parentNode.insertBefore(a,m)})(document,523)</script>