Documentation Site
You can view the documentation here: https://butchland.github.io/fastai_xla_extensions
pip install fastai_xla_extensions
Configure TPU Environment Access
The Pytorch XLA package requires an environment supporting TPUs (Kaggle kernels, GCP or Colab environments required).
Nominally, Pytorch XLA also supports GPUs so please see the Pytorch XLA site for more instructions.
If running on Colab, make sure the Runtime Type is set to TPU.
#colab
# install pytorch 1.7.1 b/c fastai doesn't support pytorch 1.8 just yet
!pip install -Uqq --no-cache-dir torch==1.7.1+cu101 torchvision==0.8.2+cu101 torchtext==0.8.0 -f https://download.pytorch.org/whl/torch_stable.html
#colab
!pip install -Uqq fastcore --upgrade
!pip install -Uqq fastai==2.3.0
This is the official way to install Pytorch-XLA 1.7 as per the instructions here
#colab
!pip install -Uqq cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.7-cp37-cp37m-linux_x86_64.whl
import warnings
try:
import torch_xla
except ImportError as e:
if DEBUG: warnings.warn('TPU environment not available')
from fastai.vision.all import *
#hide_output
from fastai_xla_extensions.all import *
Single and Multi Core TPU Modes
The fastai_xla_extensions
package allows you to train your models using either a single TPU core or multiple TPU cores.
Using multi TPU cores is usually faster than single TPU core, but due to some limitations, not all fastai features are supported in multi TPU core mode.
Multi Core TPU mode
In order to run in multi TPU core mode, the
fastai_xla_extensions
package patched some additional methods to theLearner
class.These additional methods are all prefixed by
xla
. For example, to run theLearner.fit
method in multi core tpu mode, you can invoke the equivalentLearner.xla_fit
method, with almost the same parameters.
Note that in multi core TPU mode, when calling normal learner or data loader methods (not prefixed by xla
), the process does not use the multi core TPUs, but instead runs them using the CPU.
List of
xla
prefixed methods- xla_fit
- xla_fine_tune
- xla_fit_one_cycle
- xla_fit_sgdr
- xla_fit_flat_cos
- xla_lr_find
- xla_get_preds
Additional
xla
Parameters (for Multi Core TPU mode)All
xla
prefixed methods support the following additional parameters:num_cores: this is defaulted to 8 (representing the number of TPU cores per TPU device) - valid values are
8
or1
on Google Colab. For running GCP TPUs, other values maybe possible. See the Pytorch XLA page for more details.start_method: this is defaulted to
fork
which the only value supported on Google Colab. On GCP, other values supported includespawn
. Again, see the Pytorch XLA page for more details.master_cbs: these are callbacks that are run only on the master ordinal process (rank = 0). This allows the addition of callbacks that may interfere with each other if run simultaneously on more than one process by restricting them to run on only one process - the master ordinal process.
Single Core TPU mode
In order to run in single TPU core mode, the
fastai_xla_extensions
package added a methodto_xla
to theLearner
class.After invoking this method, the learner will now use an associated TPU core device to the learner process which will be used when training the model.
This means we can use the same
fit
methods inLearner
that we normally use to train the model when using a GPU or a CPU.Also note that the single TPU core mode is incompatible with the Multi TPU Core mode, and that the kernel will have to restarted once the single TPU core mode is set. This is because once a TPU core is associated with the main process, it cannot be unset without restarting the kernel.
Multi TPU Core Example
Build a MNIST classifier -- adapted from fastai course Lesson 4 notebook
Load MNIST dataset
path = untar_data(URLs.MNIST_TINY)
Create Fastai DataBlock
datablock = DataBlock(
blocks=(ImageBlock,CategoryBlock),
get_items=get_image_files,
splitter=GrandparentSplitter(),
get_y=parent_label,
item_tfms=Resize(28),
# batch transforms run on the CPU and are slow, so only using the Normalize batch transform
batch_tfms=[Normalize.from_stats(*imagenet_stats)]
)
datablock.summary(path)
Create the dataloader
dls = datablock.dataloaders(path, bs=8) # use a small batch size as mnist_tiny dataset is very small
dls.show_batch()
Create a Fastai CNN Learner
learner = cnn_learner(dls, resnet18, metrics=accuracy, concat_pool=False) # fastai AdaptivePooling causes a lowering which slows down the training so using nn.AvgPool2D
learner.summary()
Using the lr_find
works
learner.xla_lr_find()
Run one cycle training.
learner.xla_fit_one_cycle(5,lr_max=slice(1e-4,0.02))
Further fine-tuning
learner.xla_fit_one_cycle(5,lr_max=slice(7e-4, 1e-3))
Model params are using CPU (since it is the spawned process models that are moved to TPU)
one_param(learner.model).device
Plot loss seems to be working fine.
learner.recorder.plot_loss()
Single TPU Core Example
Build a MNIST classifier -- adapted from fastai course Lesson 4 notebook
Load MNIST dataset
path = untar_data(URLs.MNIST_TINY)
Create Fastai DataBlock
datablock = DataBlock(
blocks=(ImageBlock,CategoryBlock),
get_items=get_image_files,
splitter=GrandparentSplitter(),
get_y=parent_label,
item_tfms=Resize(28),
# affine transforms are performed on the CPU, other batch transforms are done on the TPU
batch_tfms=aug_transforms(do_flip=False,min_scale=0.8)
)
datablock.summary(path)
Create the dataloader
dls = datablock.dataloaders(path)
dls.show_batch()
Create a Fastai CNN Learner
learner = cnn_learner(dls, resnet18, metrics=accuracy)
learner.summary()
Set Learner to XLA mode This will setup the learner to use the XLA Device
learner.to_xla()
Using the lr_find
works
learner.lr_find()
Run one cycle training.
learner.fit_one_cycle(5,lr_max=slice(1e-4,0.02))
Further fine-tuning
learner.fit_one_cycle(5,slice(7e-4, 1e-3))
Model params are using TPU
one_param(learner.model).device
Dataloader device is set to None (as it is the device mover transform that moves the input to the TPU)
learner.dls.device is None
Plot loss seems to be working fine.
learner.recorder.plot_loss()
from fastai_xla_extensions.utils import print_aten_ops
print_aten_ops()
Samples
Other examples of fastai notebooks using the fastai_xla_extensions package are also available here:
- Vision (Multiple TPU Core) - CIFAR
Vision (Multiple TPU Core) - Plain pytorch training with minimal fastai
Kaggle Notebook version of Vision (Multiple TPU Core) - Pets
NEW ** Vision (Multi Core) using Pytorch Image Models (timm models) with Albumentations
More samples will be added in the future as we fix issues and implement more capabilities.
Status
The fastai XLA extensions library is still in early development phase but is getting close to a 1.0 release.
The Multi Core TPU mode has been well tested on vision tasks (Classification) but the other modules (Tabular, Collab Filtering, and Text) have not.
If you use the package and encounter issues, please file bugs and issues encountered in the project's issue tracker.
Feel free to suggest enhancements or feature requests. (Though they might not necessarily be addressed any time soon).
If you wish to contribute to the project, fork it and make pull request.
This project uses nbdev -- a jupyter notebook first development environment and is being developed on Colab.
Development Setup
see this note on how to setup your environment for development