将YoloV3模型的pytorch权重文件转化为onxx

pytorch训练有两种方式保存权重文件，一是保存模型本身文件另一种是保存state_dict文件
保存state_dict文件速度较快但要先读取模型

torch.save(model.state_dict(), 'logs/Epoch%d-Total_Loss%.4f-Val_Loss%.4f.pth'%((epoch+1),total_loss/(epoch_size+1),val_loss/(epoch_size_val+1)))

torch.save(model, 'logs/Epoch%d-Total_Loss%.4f-Val_Loss%.4f.pt' % ((epoch + 1), total_loss / (epoch_size + 1), val_loss / (epoch_size_val + 1)))

state_dict

参考yolov3中训练文件，其中YoloBody是训练中的模型代码，config是定义的类名和类数量等数据，会影响最后的推理输出层，在这里需要

import torch
from nets.yolo3 import YoloBody
from utils.config import Config

config=Config
        # ---------------------------------------------------#
        #   建立yolov3模型
        # ---------------------------------------------------#
net = YoloBody(config)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

state_dict= torch.load('Epoch259-Total_Loss37.5095-Val_Loss38.6587.pth', map_location=device)
net.load_state_dict(state_dict)
net.eval()
net.cuda()
batch_size = 1  #批处理大小
input_shape = (3,416,416)   #输入数据

input_data_shape = torch.randn(batch_size, *input_shape, device=device)

torch.onnx.export(net, input_data_shape, "helmet.onnx", verbose=True)

model

torch.load读取的直接为net也不需要导入模型代码

import torch



device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

net= torch.load('Epoch5-Total_Loss90.9717-Val_Loss83.7143.pt', map_location=device)
net.eval()
net.cuda()
batch_size = 1  #批处理大小
input_shape = (3,416,416)   #输入数据

input_data_shape = torch.randn(batch_size, *input_shape, device=device)

torch.onnx.export(net, input_data_shape, "helmet2.onnx", verbose=True)