Using IMatrixMultiplyLayer for batch matrix multiplication

Description

I am using torch2trt to convert a pytorch model to tensorrt, but my model has several unsupported methods which I need to convert myself. I managed to do it except for one function that I’m stuck at: torch.bmm(). This performs a batch matrix-matrix product of matrices, i.e. inputs of shape (b*n*m) and (b*m*p) give output of shape (b*n*p).

I managed to write a working converter for torch.matmul() using the IMatrixMultiplyLayer. I used the module tester in torch2trt to confirm that the error is zero when I gave it inputs of size (1*n*m) and (1*m*p). From the documentation, I expect the converter to also work for batched inputs of size (b*n*m) and (b*m*p). But when compared with the output of torch.bmm(), it gives a significant error. Am I not understanding the documentation correctly? Does IMatrixMultiplyLayer not work like torch.bmm() ?

Environment

TensorRT Version: 5.1.5.0
GPU Type: GeForce GTX 1080
Nvidia Driver Version: 410.48
CUDA Version: 10.0
CUDNN Version: 7.6.4
Operating System + Version: ubuntu 16.04 LTS
Python Version: 2.7
PyTorch Version: 1.4.0

Relevant Files

My modified version torch2trt can be found here. The model I’m trying to convert to is PointNet, though I had to adjust the definition a little to make it work. Running it does not give any error, and on cursory glance the output doesn’t look obviously wrong.

Steps To Reproduce

  • git clone https://github.com/emilyfy/torch2trt
  • cd torch2trt
  • python -m torch2trt.test
  • Error is seen on the last line, e.g.
    torch2trt.converters.matmul.test_batchmatmul_basic | float32 | [(10, 3, 4), (10, 4, 5)] | {‘max_batch_size’: 10} | 1.13E+01 | 6.31e+04 | 1.81e+04 | 0.0795 | 0.113 |

The output should be small error for torch2trt.converters.matmul.test_matmul_basic and huge error in test_batchmatmul_basic(). The output of the run is also given here.

edit: figured out the problem. There is nothing wrong with IMatrixMultiplyLayer, it was just that torch2trt’s test code didn’t expect the output to be batched so the error was all computed with respect to the first batch.

Hi @emilyfy,

For issues with torch2trt, please file an issue on Issues · NVIDIA-AI-IOT/torch2trt · GitHub

HI @NVES_R, the issue I’m having is not with torch2trt. It’s working fine, it just didn’t support an operation I want and I wrote it myself. The issue is with tensorrt’s IMatrixMultiplyLayer, it’s not giving the result I expected. torch2trt is merely acting as my testing platform.

HI @NVES_R, I figured out the problem. There is nothing wrong with IMatrixMultiplyLayer, it was just that torch2trt’s test code didn’t expect the output to be batched so the error was all computed with respect to the first batch. Thank you and sorry for the trouble.

1 Like