晟腾 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

‍