1. x = torch.randn(1, 3, 416, 416)
    2. input_names = ["input"]
    3. out_names = ["output"]
    4. INPUT_DICT = ""#pth模型路径
    5. OUT_ONNX = ""#onnx模型路径
    6. #定义网络结构
    7. model = Network()
    8. #加载参数
    9. model_dict = model.state_dict()
    10. pretrained_dict = torch.load(INPUT_DICT, map_location=torch.device('cpu'))
    11. pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) ==  np.shape(v)}
    12. model_dict.update(pretrained_dict)
    13. model.load_state_dict(model_dict)
    14. #gpu
    15. # if torch.cuda.device_count() > 1:
    16. #     print("let's use", torch.cuda.device_count(), "Gpus!")
    17. #     model = DataParallel(model)
    18. model.eval()
    20. torch.onnx.export(model, x, OUT_ONNX, export_params=True, training=False, input_names=input_names, output_names=out_names, opset_version=11,do_constant_folding=False)
    21. #print('please run: python3 -m onnxsim test.onnx ?test_sim.onnx\n')#??????????????????????????????????????
    22. print('convert done!\n')