torch转onnx

最新推荐文章于 2024-05-14 20:58:16 发布

junkangchen

最新推荐文章于 2024-05-14 20:58:16 发布

阅读量208

点赞数

文章标签：深度学习人工智能 python

本文链接：https://blog.csdn.net/cjk_cynosure/article/details/127759907

版权

import torchvision

# pretrained的话是没有参数的
# resnet34 是继承torch.nn.Module的模型
resnet34 = torchvision.models.resnet34(pretrained=True)

# 导出到onnx

input = torch.randn(1,3, 224, 224)
# batch=1
torch.onnx.export(resnet34, input,"resnet34.onnx", verbose=True, opset_version=11, export_params=True)
# 动态batch
torch.onnx.export(resnet34, input,"resnet34.onnx", verbose=True, opset_version=11, export_params=True,input_names = ["input"],output_names = ["output"],dynamic_axes = {'input' : {0 : 'batch_size'},'output' : {0 : 'batch_size'} })