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:定义输出层名