Good evening,
When using torch.bmm() to multiply many (>10k) small 3x3 matrices, we hit a performance bottleneck apparently due to cuBLAS heuristics when choosing which kernel to call. For example, the colab notebook below shows that for 2^15 matrices the call takes 2s but only 0.5s for 2^16 matrices.
What’s the easiest way to fix this, keeping in mind that we’d like to keep the differentiability via autograd ?
Minimal Example Colaboratory: https://colab.research.google.com/drive/1BsqH2xYe61gzC8C4YuGCGmnCUzhSwBcS 15
StackOverflow discussion on batched GEMMs: https://stackoverflow.com/questions/48519861/strange-cublas-gemm-batched-performance 17
Thanks,
Séb