目录
1.更改模型框架
2.重新训练模型
3.将重新训练的模型转换为onnx模型
4.将onnx模型转为rknn模型
5.pytorch模型保存注意事项
Yolov5官网模型训练框架为pytorch1.7.0,并且模型框架中部分算子在RKNN_Toolkit_V1.6.0中是不支持的,因此常规的模型转换过程是无法进行模型转换的。这里我们做一个详细的说明。
1.更改模型框架
更改模型框架,也就是替换掉不支持的模型算子,参考https://github.com/airockchip/yolov5/ 下面的代码,主要是对models/commom.py 进行以下修改:
第30行 将Conv定义中的激活函数进行修改self.act = nn.Hardswish() if act else nn.Identity() 替换为self.act = nn.ReLU() if act else nn.Identity() 第62行 将self.act = nn.LeakyReLU(0.1, inplace=True)替换为self.act = nn.ReLU() 第104行,将self.conv = Conv(c1 * 4, c2, k, s, p, g, act) 替换为self.conv = Conv(c1, c2, k, 2, p, g, act) 第107行 将return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))替换为return self.conv(x)
2.重新训练模型
参考 train.py ,需要更换算法模型和训练用的图片和标签。
3.将重新训练的模型转换为onnx模型
转换onnx参考:2onnx.sh。
4.将onnx模型转为rknn模型
转换rknn参考:rknn/onnx2rknn.py。
注意要修改rknn/onnx2rknn.py里面的target_platform。
if __name__ == '__main__':
# Create RKNN object
rknn = RKNN(verbose=True)
# pre-process config
print('--> config model')
rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], reorder_channel='0 1 2', target_platform=['rv1126'], output_optimize=1)
print('done')
5.pytorch模型保存注意事项
在使用 load_pytorch 接口加载 pytorch 模型时, pytorch 模型只 能 用 torch.jit.trace 保 存 的 模 型。
用torch.save(net.state_dict())保存的模型,保存的只是模型的参数,并没有网络结构。只凭一个网络参数, RKNN Toolkit 无法构建相应的网络。
说明:更多模型转换注意事项见《Rockchip_Trouble_Shooting_RKNN_Toolkit_V1.6.0_CN.pdf》。
参考文献:
https://github.com/airockchip/yolov5
https://github.com/rockchip-linux/rknn-toolkit/tree/master/doc



