Encountered known unsupported method torch.bmm

[Water2style]

begin to convert:
Warning: Encountered known unsupported method torch.bmm
[TensorRT] ERROR: INVALID_ARGUMENT: Cannot find binding of given name: input_0
tensor(2.24113e-05, device='cuda:0', grad_fn=)

[jaybdub]

Hi Water2style,

Thanks for reaching out!

Currently, there is not a registered converter for the torch.bmm operation. It may be possible to implement the converter yourself as described in the README. The TensorRT python documentation may be found here

https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Graph/Layers.html#imatrixmultiplylayer

Please let me know if this helps, or you run into any issues. I'm curious if you're able to get this to work.

Best,
John

[Water2style]

Hi Water2style,

Thanks for reaching out!

Currently, there is not a registered converter for the torch.bmm operation. It may be possible to implement the converter yourself as described in the README. The TensorRT python documentation may be found here

https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Graph/Layers.html#imatrixmultiplylayer

Please let me know if this helps, or you run into any issues. I'm curious if you're able to get this to work.

Best,
John

Thank you john!

[Water2style]

Hi Water2style,

Thanks for reaching out!

Currently, there is not a registered converter for the torch.bmm operation. It may be possible to implement the converter yourself as described in the README. The TensorRT python documentation may be found here

https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Graph/Layers.html#imatrixmultiplylayer

Please let me know if this helps, or you run into any issues. I'm curious if you're able to get this to work.

Best,
John

Hi John, it works! Thank you

@tensorrt_converter('torch.bmm')
def convert_bmm(ctx):
input_a = ctx.method_args[0]
input_b = ctx.method_args[1]
input_a_trt, input_b_trt = trt_(ctx.network, input_a, input_b)
output = ctx.method_return
layer = ctx.network.add_matrix_multiply(input_a_trt,input_b_trt)
output._trt = layer.get_output(0)

Other code use your README,
such as: x = torch.ones((1, 3, 224, 224)).cuda()

And the results: check the output against Pytorch:
tensor( 2.5........-05 , .......)
I will test the model in my image data soon

[jaybdub]

Great to hear! Let me know how it works with your image data.

Best,
John

[Water2style]

Great to hear! Let me know how it works with your image data.

Best,
John

It works well with my image data !
You can add this converter to "torch2trt/torch2trt/converters/ ” in this git. LOL

[summer-wind-hi]

@jaybdub hi John, I am a first learner of tensorrt, now I want to define and use a custom layer by tensorrt_converter, is there any tutorials or demos, especially how to implement a layer such as layer = ctx.network.add_pooling ? I need to transfer x[x == 0] = e-3 from pytorch model to trt model. Thanks.

[tuqingyun]

@Water2style
Do you have solved the bmm by add th code above problem?

[a1270107629]

@Water2style I add this converter to "torch2trt/torch2trt/converters/ ” in this git,but it cann't work.