Deep learning format conversion notes (1): Detailed explanation of converting model file pt to onnx format

Common model files include model files with suffixes of .pt, .pth, .pkl, and these types of model files are not different in format but with different suffixes. To save model files is often torch.save(), The suffix is ​​different simply because everyone's preferences are different. Usually pth and pt are used.
Save :
orch.save(model.state_dict(), mymodel.pth)#Only save model weight parameters, not model structure

Call :
model = My_model(*args, **kwargs)
#here you need to re-model the structure, pthfile = r'absolute path'
loaded_model = torch.load(pthfile, map_location='cpu')
model.load_state_dict(loaded_model['model' ])
model.eval( ) #Do not enable BatchNormalization and Dropout, do not change the weight

turn

from nn.mobilenetv3 import mobilenetv3_large,mobilenetv3_large_full,mobilenetv3_small
import torch
from nn.models import DarknetWithShh
from hyp import hyp

def convert_onnx():
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model_path = 'weights/mbv3_large_75_light_final.pt' #这是我们要转换的模型
    backone = mobilenetv3_large(width_mult=0.75)#mobilenetv3_small()  mobilenetv3_small(width_mult=0.75)  mobilenetv3_large(width_mult=0.75)
    model = DarknetWithShh(backone, hyp,light_head=True).to(device)

    model.load_state_dict(torch.load(model_path, map_location=device)['model'])

    model.to(device)
    model.eval()
    dummy_input = torch.randn(1, 3, 32, 32).to(device)#输入大小   #data type nchw
    onnx_path = 'weights/mbv3_large_75_light_final.onnx'
    torch.onnx.export(model, dummy_input, onnx_path, input_names=['input'], output_names=['output'],opset_version=11)
    print('convert retinaface to onnx finish!!!')

if __name__ == "__main__" :
    convert_onnx()

Conversion result