pytorch模型转化为onnx

import torch, onnx, collections

print('notice !!!! ----> use python3 run this script!!! \n')
INPUT_DICT = 'res/model_test/eye_fbresnet10_motion.pth'//input path
OUT_ONNX = 'res/model_test/eye_fbresnet10_motion.onnx'//output path

x = torch.randn(1, 3, 224, 224);
input_names = ["input"];
out_names = ["output"];

xmodel= torch.load(INPUT_DICT, map_location=torch.device('cpu'))
xmodel.eval()

torch.onnx.export(xmodel, x, OUT_ONNX, export_params=True, training=False, input_names=input_names, output_names=out_names)
print('please run: python3 -m onnxsim test.onnx  test_sim.onnx\n')//对转化后的模型，运行这个命令，简化符号
print('convert done!\n')