def forward(self, x):
results = []
for i in range(self.nl):
dfl = self.cv2[i](x[i]).permute(0, 2, 3, 1).contiguous()
cls = self.cv3[i](x[i]).permute(0, 2, 3, 1).contiguous()
results.append(torch.cat([cls, dfl], -1))
return tuple(results)
新建export.py文件
#encoding:utf-8
from ultralytics import YOLO
import onnx
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default='./weights/yolov8n.pt', help='initial weights path')
#================================================================
opt = parser.parse_args()
print(opt)
model_path = opt.weights
# load a pretrained model
model = YOLO(model_path)
# export onnx
success = model.export(format='onnx', opset=11, simplify=True, dynamic=False, imgsz=640)
assert success
model = onnx.load(model_path.replace(".pt",".onnx"))
# 修改输入输出张量的名称
idx_start = 0
for input in model.graph.input:
for node in model.graph.node:
# 如果当前节点的输入名称与待修改的名称相同,则将其替换为新名称
for i, name in enumerate(node.input):
if name == input.name:
node.input[i] = "images"
input.name = "images"
idx_start += 1
idx_start = 0
for output in model.graph.output:
for node in model.graph.node:
# 如果当前节点的输入名称与待修改的名称相同,则将其替换为新名称
for i, name in enumerate(node.output):
if name == output.name:
node.output[i] = "output" + str(idx_start)
output.name = "output" + str(idx_start)
idx_start += 1
# 保存修改后的模型
onnx.save(model, model_path.replace(".pt",".onnx"))