赞
踩
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)
#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=12, 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] = "data" input.name = "data" 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] = "out" + str(idx_start) output.name = "out" + str(idx_start) idx_start += 1 # 保存修改后的模型 onnx.save(model, model_path.replace(".pt",".onnx"))
Namespace(weights='./weights/yolov8n.pt')
finished exporting onnx
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。