xla
xla copied to clipboard
pytorch
How to load a gpu trained model on TPU for evaluation
❓ Questions and Help
Hello, I am loading a GPU trained model on map_location=cpu and then doing "model.to(device)" where device is xm.xla_device(n=device_num,devkind="TPU") but on testing the cpu processing time and the tpu processing time is the same. Please let me know what I can do about it.
Thank you
I am not sure if I fully get your question. Do you means using TPU for eval does not provide any speed up compared to CPU?
That's right, I trained a model on GPU device and now I am using the model to evaluate images on TPU device but the processing time is worse than that of using a CPU device.
You might want to start with this doc. My guess is that the model you are trying to evaluate is fairly small or you might be keep recompiling the graph due to different input shape etc. Feel free to check out our resnet example in https://github.com/pytorch/xla/blob/master/test/test_train_mp_imagenet.py where we only compile once and execute for many epochs.