如何在onnx中知道模型的输入大小?

时间:2019-07-15 02:46:49

标签: pytorch onnx

输入的大小未在pytorch中指定。只需调整内核大小即可生成输出。 WinMLDashboard显示图像输入的宽度和高度。那怎么可能?

1 个答案:

答案 0 :(得分:1)

当您将网络从pytorch序列化到onnx时,您的意思是?因为从pytorch导出时,您需要根据documentation

定义输入的大小
dummy_input = torch.randn(10, 3, 224, 224, device='cuda')
model = torchvision.models.alexnet(pretrained=True).cuda()
input_names = [ "actual_input_1" ] 
output_names = [ "output1" ]

torch.onnx.export(model, dummy_input, "alexnet.onnx", verbose=True, 
input_names=input_names, output_names=output_names)