当前位置: 代码迷 >> 综合 >> PyTorch模型转换为ONNX格式模型
  详细解决方案

PyTorch模型转换为ONNX格式模型

热度:88   发布时间:2024-02-25 02:28:43.0

测试PyTorch版本:1.2.0

测试代码:

import torch.onnxmodel        = ******** # model类型:torch.nn.Module
dummy_input  = torch.randn(1, 3, 224, 224) # 模型输入维度为(1, 3, 224, 224)
input_names  = ["input"]
output_names = ["output"]
torch.onnx.export(model, dummy_input, "model_onnx.onnx", verbose=True,input_names=input_names, output_names=output_names)

 

  相关解决方案