Fastai v2 TPU support

[ilovescience (Tanishq)]

@butchland

As for some monkey-patching, I had to do some of that (especially on getting a default_device to return a TPU) because I don’t think Sylvain or Jeremy initially considered an environment where if a GPU is not available, that the default would be anything other than a CPU… (which would be the case where a TPU was available)

You’ll probably have to ask Jeremy about that. Personally, I don’t think this is something that’s strictly necessary, but it’s likely a decision that Jeremy will need to make about whether or not to include TPUs as a default. But I guess I didn’t really have a problem with that and instead was talking about the separate optimizer classes you created, which is not necessary.

@butchland

In any case, our goal is to make it so that using a TPU on fastai would require minimal changes to your existing fastai notebooks or code.

Yep, my goal the is same!

@kcturgutlu

I previously attempted to created a similar Learner class in here for fastai-v1 to work with mutlicore TPU.

I see you are using my kernel (developed with the help of the PyTorch XLA team and @abhi1thakur) .
Which version of the kernel is the working one? The latest one is just a quick save, and while there is an older working one, I am not sure if that’s your final fastai version or if there’s more to it?

@kcturgutlu

I think it would be more important to get multicore working since most TPU devices offered publicly (Kaggle, Colab) are of that kind and it would allow us to use TPUs for the main reason - speed.

Exactly my thoughts! I have tried single-core TPU training with very little benefit. Hence, I have been focusing on multi-core TPU training.

Anyway, I will work on it more today and keep you guys updated in this thread!