PyTorch Slow Batch matrix multiplication on GPU
I am using Batch Matrix Multiplication on 2 3d tensors of sizes (100 , 128 , 128 ) each.

import torch
a = torch.randn(100,128,128)
b = torch.randn(100,128,128)

import time

t0 = time.time()
torch.bmm(a,b)
print(time.time() - t0)



  0.03233695030212402


Now if i do the same thing on GPU it takes a lot longer

a = a.cuda()
b = b.cuda()
t0 = time.time()
torch.bmm(a,b)
print(time.time() - t0)



  30.574532985687256


Why does it take so long to solve on GPU?

I have a GTX 1050 Ti 4GB

And processor core i3-6100 3.7Ghz

GPU:

30.57 secs is the total time taken by following steps:


CPU launches kernels* on the device (GPU)
CPU allocates memory on GPU
CPU copies input data to GPU
CPU launches kernels on GPU to process the input data
CPU copies output results to itself


*Kernel is a serial code which is a small part of the original code.

CPU

Whereas, 0.0323 secs is the total time taken by either:


Intra-CPU communications
CPU and main memory communications


These are very fast and they also have to execute fewer commands to get work done.

Conclusion

Hence it's a trade-off between the complexity of the commands (CPU commands being more complex) and speed-up due to parellelism and simpler nature of commands (GPU can fire more parallel threads and they are much simpler in nature).

If there were more computations to be done, then GPU speed-up would overwhelm the extra time taken for CPU-GPU communications. 
