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'
引用 皆我百晓生 小程序回复内容作答:
这个错误是由于在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脚本就不会报错了。