Notas de conversión de formato de aprendizaje profundo (1): explicación detallada de la conversión del archivo de modelo pt a formato onnx

Los archivos de modelo comunes incluyen archivos de modelo con sufijos de .pt, .pth, .pkl, y estos tipos de archivos de modelo no son diferentes en formato pero con diferentes sufijos. Para guardar archivos de modelo a menudo se utiliza torch.save (), El sufijo es diferente simplemente porque las preferencias de todos son diferentes. Normalmente se utilizan pth y pt.
Guardar :
orch.save (model.state_dict (), mymodel.pth) # Solo guarda los parámetros de peso del modelo, no la estructura del modelo

Llame a :
model = My_model (* args, ** kwargs)
#aquí debe volver a modelar la estructura, pthfile = r'absolute
path'loaded_model = torch.load (pthfile, map_location = 'cpu')
model.load_state_dict (loaded_model ['modelo' ])
model.eval ( ) #No habilite BatchNormalization y Dropout, no cambie el peso

giro

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()

Resultado de conversión