1、需要修改,ultralytics\nn\modules\block.py
class C2f
def forward(self, x):
"""Forward pass through C2f layer."""
"""Forward pass of a YOLOv5 CSPDarknet backbone layer."""
y = list(self.cv1(x).chunk(2, 1))
y.extend(m(y[-1]) for m in self.m)
return self.cv2(torch.cat(y, 1))
# '''转onnx模型时修改'''
# x = self.cv1(x)
# x = [x, x[:, self.c:, ...]]
# x.extend(m(x[-1]) for m in self.m)
# x.pop(1)
# return self.cv2(torch.cat(x, 1))
2.需要修改,ultralytics\nn\modules\head.py
class Detect
def forward(self, x):
"""Concatenates and returns predicted bounding boxes and class probabilities."""
shape = x[0].shape # BCHW
for i in range(self.nl):
x[i] = torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)
if self.training:
return x
elif self.dynamic or self.shape != shape:
self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5))
self.shape = shape
# x_cat = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2)
# if self.export and self.format in ('saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs'): # avoid TF FlexSplitV ops
# box = x_cat[:, :self.reg_max * 4]
# cls = x_cat[:, self.reg_max * 4:]
# else:
# box, cls = x_cat.split((self.reg_max * 4, self.nc), 1)
# dbox = dist2bbox(self.dfl(box), self.anchors.unsqueeze(0), xywh=True, dim=1) * self.strides
# y = torch.cat((dbox, cls.sigmoid()), 1)
# return y if self.export else (y, x)
pred = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2).permute(0, 2, 1)
return pred
修改后,模型不需要重新训练,运行以下,导出即可,其余步骤不变,再次训练和识别的时候改回原版,否则报错
from ultralytics import YOLO
# Load a model
model = YOLO('path/to/best.pt') # load a custom trained
# Use the model
success = model.export(format="onnx",imgsz=320,half = True,optimize =True)