Xla: Training on TPU is slower than GPU

[ibeltagy]

Thank you for all the effort to run pytorch with TPUs.

I tried the simple MNIST example with the prebuilt docker image and it works fine. However, it seems that running it on a TPU is slower than GPU (numbers below). Is this a known issue or is there something I am doing wrong?

Titan V:  5.13second/epoch
TPU v3-8: 8.18second/epoch

[dlibenzi]

Which command are you running for TPU?

[ibeltagy]

I am just running the mnist example
docker run --shm-size 16G -e XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470" gcr.io/tpu-pytorch/xla:r0.1 python /pytorch/xla/test/test_train_mnist.py

[dlibenzi]

You can try adding --batch_size=256 (or 512).
But also, how do you measure EPOCH time?

[ibeltagy]

Problem fixed. Turns out I used time.process_time() instead of time.time() which is not what we want to measure. I also changed number of workers reading the dataset from 4 to 2 which impacted the performance a lot.
The numbers I am getting right now are:

Titan V:  10.01second/epoch
TPU v3-8: 4.88second/epoch

I will close the issue. Thanks.

[dlibenzi]

Also, MNIST training EPOCH is so short, that does not make much sense to benchmark that, as we have a fixed cost in XLA/TPU which does not get amortized.
Testing test_train_imagenet.py or even test_train_cifar.py is a better compare.
One more thing is that a TPU v3 has 8 cores, so performance should be considerably better than a single Titan.