PyTorch模型转换为ONNX格式模型

最新推荐文章于 2024-08-13 08:16:12 发布

gy笨瓜

最新推荐文章于 2024-08-13 08:16:12 发布

阅读量2.1k

点赞数 2

文章标签： pytorch ONNX

本文链接：https://blog.csdn.net/u012633319/article/details/108913301

版权

测试PyTorch版本：1.2.0

测试代码：

import torch.onnx

model        = ******** # 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)