训练好的.pth模型转为.onnx模型

#根据自己的模型做修改，我的训练模型是根据Bubbliiiing，b导的YOLO系列
#在此之前要安装onnx
#conda install onnx
import torch.onnx
from nets.yolo import YoloBody # 该文件为模型结构文件
#下面路径按照自己的模型路径改
pt_model_path = ‘./logs/best_epoch_weights.pth’ #.pth模型路径
onnx_model_path = ‘./best.onnx’ #生成的.onnx模型路径
anchors_mask = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]
model = YoloBody(anchors_mask,num_classes = 2,pretrained=False) #num_classes为你的训练的样本类别数
device = torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’)
model.load_state_dict(torch.load(pt_model_path, map_location=device))
input_tensor = torch.randn(1, 3, 224, 224)
input_names = [‘input’]
output_names = [‘output’]
torch.onnx.export(model, input_tensor, onnx_model_path, verbose=True, input_names=input_names, output_names=output_names)