pytorch轉onnx


一個輸入:

    input_tensor = torch.randn([1, 3, 256, 512])
    print ("Exporting to ONNX: ", onnx_save_name)
    torch_onnx_out = torch.onnx.export(model, input_tensor, onnx_save_name, 
                        export_params=True,
                        verbose=True, 
                        input_names=['label'],
                        output_names=["synthesized"],
                        opset_version=11)

多個輸入:

    input_tensor = torch.randn([1, 3, 256, 512])
    mask_tensor = torch.randn([1, 3, 256, 512])
    print ("Exporting to ONNX: ", onnx_save_name)
    torch_onnx_out = torch.onnx.export(model, (input_tensor,mask_tensor), onnx_save_name, 
                        export_params=True,
                        verbose=True, 
                        input_names=['label','mask'],
                        output_names=["synthesized"],
                        opset_version=11)

 


免責聲明!

本站轉載的文章為個人學習借鑒使用,本站對版權不負任何法律責任。如果侵犯了您的隱私權益,請聯系本站郵箱yoyou2525@163.com刪除。



 
粵ICP備18138465號   © 2018-2025 CODEPRJ.COM