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.