华为计算微信公众号
昇腾AI开发者公众号
华为计算微博
华为计算今日头条
已经训练好的离线模型，要将它转换成onnx格式，batchsize怎么设置（训练和验证时，batchsize=2，测试时我想设置为1），下面代码我设置的是2
# pytorch 模型转换为 onnx 格式 import torch import torchvision dummy_input = torch.randn(2, 3, 224, 224) model = torch.load('D:/Project/Project3/utils/network_vgg16.pth') # 给输入输出取个名字 input_names = ["input_1"] output_names = ["output_1"] torch.onnx.export(model, dummy_input, "vgg16.onnx", verbose=True, input_names=input_names, output_names=output_names)
我要发帖子
点赞
收藏
回复
分享
举报
已解决 已解决
已经训练好的离线模型，要将它转换成onnx格式，batchsize怎么设置（训练和验证时，batchsize=2，测试时我想设置为1），下面代码我设置的是2
# pytorch 模型转换为 onnx 格式 import torch import torchvision dummy_input = torch.randn(2, 3, 224, 224) model = torch.load('D:/Project/Project3/utils/network_vgg16.pth') # 给输入输出取个名字 input_names = ["input_1"] output_names = ["output_1"] torch.onnx.export(model, dummy_input, "vgg16.onnx", verbose=True, input_names=input_names, output_names=output_names)