晟腾 YOLOV5 转onnx
晟腾 pytorch 转onnx
以下操作都是基于官方yolov5s.pt 为案例
1 修改yolov5 export.py 文件
/YOLOv5/yolov5/models/common.py 中 Foucs 类中的 forward方法。
修改yolov5 网络的 Focus 结构。添加逻辑,当作为onnx 导出时修改Foucus 的结构。
def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2)
if torch.onnx.is_in_onnx_export():
a, b = x[..., ::2, :].transpose(-2, -1), x[..., 1::2, :].transpose(-2, -1)
c = torch.cat([a[..., ::2, :], b[..., ::2, :], a[..., 1::2, :], b[..., 1::2, :]], 1).transpose(-2, -1)
return self.conv(c)
else:
return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))
说明 torch的transpose方法 是交换两个维度。假如x.shape=(1,3,320,640)
x.transpose(-2,-1)---->shape(1,3,640,320) 类似于转置操作。
下图为修改之前的Focus 操作图表示。
DeepinScreenshot_select-area_20230104153056
下图为修改后的图表示focus
DeepinScreenshot_select-area_20230104153037
2 修改 yolo.py 文件
主要目的是控制转换成onnx 的 后处理部分,也就是还原成三个输出的格式:
1x255x80x80; 1x255x40x40; 1x255x20x20;
class Detect(nn.Module):
def forward(self, x):
for i in range(self.nl):
x[i] = self.mi # conv
# ******* 修改内容 *****
if torch.onnx.is_in_onnx_export():
continue
# **********************
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()
原始的yolov5 会将三个输出进一步整理格式(1x3x85x20x20)
修改 onnx 的算子版本为11
修改 models.export.py 文件。
torch.onnx.export(model, img, f, verbose=False, opset_version=11, input_names=['images'],
output_names=['classes', 'boxes'] if y is None else ['output'],
dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # size(1,3,640,640)
'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None)
3 最后 使用onnx-simplifer工具进行简化
此工具主要用于简化onnx的算子,比如算子融合等。
python3 -m onnxsim --skip-optimization yolov5s.onnx yolov5s_sim.onnx