yolov5中pt文件转onnx文件报错

Question

yolov5-7.0中.pt文件转onnx文件时候修改yolo.py文件中修改class Detect中的forward函数为
def forward(self, x):
z = [] # inference output
for i in range(self.nl):
x[i] = self.mi # conv
return x
后用export.py脚本转onnx模型报错：AttributeError: 'list' object has no attribute 'shape'

Answer 1

引用 皆我百晓生 小程序回复内容作答：

这个错误是由于在yolo.py文件中修改Detect类的forward函数后，返回的是一个list对象，而在export.py脚本中调用torch.onnx.export函数时，需要传入一个tensor对象作为输入。因此，需要将返回的list对象转换为tensor对象。

以下是修改的代码示例：

import torch

class Detect(torch.nn.Module):
    def __init__(self, ...):
        ...
  
    def forward(self, x):
        z = [] # inference output
        for i in range(self.nl):
            x[i] = self.mi # conv
            z.append(x[i])
        
        return torch.cat(z, dim=1) # Convert list to tensor

# 代码中使用到的...表示省略的其他代码，需要根据实际情况进行补充

在forward函数中，将z列表中的每个元素添加到新的z列表中，然后使用torch.cat函数将z列表中的tensor对象在维度1上进行拼接，得到一个新的tensor对象，作为forward函数的返回值。

修改后，再次运行export.py脚本就不会报错了。