pytorch_onnx.py

from torchreid import models
print(models.show_avai_models())

from torch.autograd import Variable
import torch.onnx
import torchvision

#H-256
#W-128
dummy_input = torch.randn(1,3,256,128)
#dummy_input = torch.randn(1,3,128,256)
model = models.build_model('resnet50',751,loss='softmax')
#model = models.resnet50mid(pretrained=True)
torch.onnx.export(model,dummy_input,"resnet50mid.onnx")