Error when converting batchnorm1d

[joechenrh]

Hi,
I'm using Python 2.7, TensorRT 5.1.5.0 and Pytorch 1.2.0 and I want to convert my model with torch2trt, but I got error in BatchNorm1d, here is the output.

[TensorRT] ERROR: Parameter check failed at: ../builder/Network.cpp::addScale::175, condition: shift.count == 0 || shift.count == weightCount
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/lib/python2.7/dist-packages/torch2trt-0.0.0-py2.7.egg/torch2trt/torch2trt.py", line 252, in torch2trt
    outputs = module(*inputs)
  File "/usr/local/lib/python2.7/dist-packages/torch/nn/modules/module.py", line 547, in __call__
    result = self.forward(*input, **kwargs)
  File "Model/models.py", line 51, in forward
    return self.layers(inputs[0])
  File "/usr/local/lib/python2.7/dist-packages/torch/nn/modules/module.py", line 547, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python2.7/dist-packages/torch/nn/modules/container.py", line 92, in forward
    input = module(input)
  File "/usr/local/lib/python2.7/dist-packages/torch/nn/modules/module.py", line 547, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python2.7/dist-packages/torch2trt-0.0.0-py2.7.egg/torch2trt/torch2trt.py", line 97, in wrapper
    converter(ctx)
  File "/usr/local/lib/python2.7/dist-packages/torch2trt-0.0.0-py2.7.egg/torch2trt/converters/BatchNorm1d.py", line 22, in convert_BatchNorm2d
    layer = ctx.network.add_shuffle(layer.get_output(0))
AttributeError: 'NoneType' object has no attribute 'get_output'

And here is my model.

>>> print(model)
Baseline(
  (layers): Sequential(
    (0): Dropout(p=0.25, inplace=False)
    (1): Linear(in_features=2000, out_features=256, bias=True)
    (2): ReLU()
    (3): BatchNorm1d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (4): Dropout(p=0.25, inplace=False)
    (5): Linear(in_features=256, out_features=128, bias=True)
    (6): ReLU()
    (7): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (8): Dropout(p=0.25, inplace=False)
    (9): Linear(in_features=128, out_features=2, bias=True)
  )
)
>>> x = torch.zeros([1, 2000]).to('cuda:1')
>>> print(model(x))
tensor([[0., 0.]], device='cuda:1', grad_fn=<AddmmBackward>)
>>> model_trt = torch2trt(model, [x])  # got error

[jaybdub]

Hi joechenrh,

Thanks for pointing this out. The BatchNorm1d converter was only able to handle 3D inputs. I have updated to add support for 2D inputs in this pull request.

#43

It has just been merged into master so you can give it a try.

Please let me know if this works for you or if you have any other questions.

Best,
John

[joechenrh]

I have pulled the latest version and it works well.
Thanks for your help, it's quite simple to use.

[jaybdub]

Glad to hear! Please let me know if you run into any other issues.

Best,
John