任务目标
- 将训练好的yolov7.pt转换为yolov7.onnx。
- 将yolov7.onnx转换为yolov7.rknn。
实验过程
说明:本次实验仅为模型转换步骤的验证,所以不涉及模型训练部分。
- 首先去官网下载源码和yolov7.pt文件。
源码地址:https://github.com/WongKinYiu/yolov7.git
模型地址:https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7.pt
-
在
models/yolo.py
中将Detect类中的前向传播过程由:
def forward(self, x):
# x = x.copy() # for profiling
z = [] # inference output
self.training |= self.export
for i in range(self.nl):
x[i] = self.m[i](x[i]) # conv
bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85)
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
if not self.training: # inference
if self.grid[i].shape[2:4] != x[i].shape[2:4]:
self.grid[i] = self._make_grid(nx, ny).to(x[i].device)
y = x[i].sigmoid()
if not torch.onnx.is_in_onnx_export():
y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
else:
xy, wh, conf = y.split((2, 2, self.nc + 1), 4) # y.tensor_split((2, 4, 5), 4) # torch 1.8.0
xy = xy * (2. * self.stride[i]) + (self.stride[i] * (self.grid[i] - 0.5)) # new xy
wh = wh ** 2 * (4 * self.anchor_grid[i].data) # new wh
y = torch.cat((xy, wh, conf), 4)
z.append(y.view(bs, -1, self.no))
if self.training:
out = x
elif self.end2end:
out = torch.cat(z, 1)
elif self.include_nms:
z = self.convert(z)
out = (z, )
elif self.concat:
out = torch.cat(z, 1)
else:
out = (torch.cat(z, 1), x)
return out
改为:
def forward(self, x):
z = [] # inference output
for i in range(self.nl):
x[i] = self.m[i](x[i])
return x
注:这一步一定要改,否则在onnx转rknn时会报错
。
修改后将yolov7.pt移动至export.py同一文件夹下,运行export.py便可以得到yolov7.onnx。
3. 去官网下载rknntoolkit2,并在
Linux
上解压,官网地址如下:
https://github.com/rockchip-linux/rknn-toolkit2
使用Miniconda创建rknn环境:
conda create -n rknn python=3.8
conda activate rknn
环境安装:
#安装python工具,ubuntu20.04默认是安装了python3.8.10
sudo apt update
sudo apt-get install python3-dev python3-pip python3.8-venv gcc
#安装相关库和软件包
sudo apt-get install libxslt1-dev zlib1g-dev libglib2.0 libsm6 \
libgl1-mesa-glx libprotobuf-dev gcc
安装rknntoolkit2:
#安装依赖库,根据rknn-toolkit2\doc\requirements_cp38-1.4.0.txt
pip3 install numpy
pip3 install -r doc/requirements_cp38-1.4.0.txt
#安装rknn_toolkit2
pip3 install packages/rknn_toolkit2-1.4.0_22dcfef4-cp38-cp38-linux_x86_64.whl
检测是否安装成功:
(.toolkit2_env) llh@-:~/project-Toolkit2$ python
Python 3.8.10 (default, Jun 22 2022, 20:18:18)
[GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> from rknn.api import RKNN
>>>
#没有报错就安装成功了
-
在rknn-toolkit2工程文件夹中浏览至./examples/onnx/yolov5,将我们转换得到的yolov7.onnx复制到该文件夹下,修改该文件夹下的test.py中的内容为自己模型的名字,要修改的地方如下:
此处我们只需要修改前两项和搭载平台的名字即可。
最后运行test.py,即可得到rknn模型。
版权声明:本文为xingman510原创文章,遵循 CC 4.0 BY-SA 版权协议,转载请附上原文出处链接和本声明。