pytorch 保存模型爲onnx

 

# An example input you would normally provide to your model's forward() method
x = torch.rand(1, 3, 512, 720).float().cuda()
input_names = ["inputs"]
output_names = ["main_out_rest", "sub1_out_rest", "sub2_out_rest", "sub4_out_rest"]
# Export the model
torch_out = torch.onnx.export(network, x, "lsc-cnn.onnx", export_params=True, verbose=True,input_names=input_names, output_names=output_names)

network :模型

x:模型輸入

"lsc-cnn.onnx" :模型保存名

export_params:是否保存參數

verbose:查看模型信息

input_name:定義輸入層名

output_name:定義輸出層名