System information
- OS Platform and Distribution : Linux Ubuntu 18.04
- ONNX version : 1.9.0
- Python version: 3.8.10
- Torch version: 1.8.0
Describe the bug
I am trying to convert torch model(.pt) to onnx format but meets unexpected error export failure: number of dims don't match in permute
, I inspect into model structure and find out where the problem is and the module is shown below:
import onnx
import torch
import torch.nn as nn
import torch.nn.functional as F
class Integral(nn.Module):
def __init__(self, reg_max= 16):
super(Integral, self).__init__()
self.reg_max = reg_max
self.register_buffer('project',
torch.linspace(0, self.reg_max, self.reg_max + 1))
def forward(self, x):
#x.shape (1, 3549, 68)
x = F.softmax(x.reshape(x.shape[0], -1, self.reg_max + 1), dim = 2)
#x.shape(1, 14196, 17)
x = F.linear(x, self.project.type_as(x)).reshape(x.shape[0],-1, 4)
return x
#model setting
model = Integral()
model.to(device) #cpu or gpu
model.eval()
#input
x = torch.randn(1, 3549, 68).to(device)
f = opt.weights.replace('.pt', '.onnx') # onnx filename
#opset_version=12
#opt.train =False
torch.onnx.export(model, x, f, verbose=False, opset_version=opt.opset_version, input_names=['images'],
training=torch.onnx.TrainingMode.TRAINING if opt.train else torch.onnx.TrainingMode.EVAL,
do_constant_folding=True,
dynamic_axes= None)
The code above gives error : RuntimeError: number of dims don't match in permute
and I don’t see any wrong with tensor operation in foward
. Can anyone help me out with this? Thanks!