Package: fastai-xla-extensions
FastAI XLA Extensions LibraryThe FastAI XLA Extensions library package allows your fastai/Pytorch models to run on TPUs using the Pytorch-XLA library.Documentation SiteYou can view the documentation here:https://butchland.github.io/fastai_xla_extensionsInstallpip install fastai_xla_extensionsHow to useConfigure TPU Environment AccessThe Pytorch XLA package requires an environment supporting TPUs (Kaggle kernels, GCP or Colab environments required).Nominally, Pytorch XLA also supports  GPUs so please see thePytorch XLA site for more instructions.If running on Colab, make sure the Runtime Type is set to TPU.Install pytorch 1.7.1Use pytorch 1.7.1 as this is latest version of pytorch supported by the fastai package#hide_output#colab# install pytorch 1.7.1 b/c fastai doesn't support pytorch 1.8 just yet!pipinstall-Uqq--no-cache-dirtorch==1.7.1+cu101torchvision==0.8.2+cu101torchtext==0.8.0-fhttps://download.pytorch.org/whl/torch_stable.htmlInstall fastaiUse the latest fastai and fastcore versions (note: due to a compatibility issue, we are using fastai==2.3.0 for now)#hide_output#colab!pipinstall-Uqqfastcore--upgrade!pipinstall-Uqqfastai==2.3.0Install Pytorch XLAThis is the official way to install Pytorch-XLA 1.7 as per theinstructions here#hide_output#colab!pipinstall-Uqqcloud-tpu-client==0.10https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.7-cp37-cp37m-linux_x86_64.whlCheck if XLA is available#colabimportwarningstry:importtorch_xlaexceptImportErrorase:ifDEBUG:warnings.warn('TPU environment not available')WARNING:root:Waiting for TPU to be start up with version pytorch-1.7...
WARNING:root:Waiting for TPU to be start up with version pytorch-1.7...
WARNING:root:TPU has started up successfully with version pytorch-1.7Import the librariesImport the fastai and fastai_xla_extensions librariesfromfastai.vision.allimport*#colab#hide_outputfromfastai_xla_extensions.allimport*Single and Multi Core TPU ModesThefastai_xla_extensionspackage 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 modeIn order to run in multi TPU core mode, thefastai_xla_extensionspackage patched some additional methods to theLearnerclass.These additional methods are all prefixed byxla. For example, to run theLearner.fitmethod in multi core tpu mode, you can invoke the equivalentLearner.xla_fitmethod, with almost the same parameters.Note that in multi core TPU mode, when calling normal learner or data loader methods (not prefixed byxla), the process does not use the multi core TPUs, but instead runs them using the CPU.List ofxlaprefixed methodsxla_fitxla_fine_tunexla_fit_one_cyclexla_fit_sgdrxla_fit_flat_cosxla_lr_findxla_get_predsAdditionalxlaParameters (for Multi Core TPU mode)Allxlaprefixed methods support the following additional parameters:num_cores: this is defaulted to 8 (representing the number of TPU cores per TPU device) - valid values are8or1on Google Colab. For running GCP TPUs, other values maybe possible. See thePytorch XLA pagefor more details.start_method: this is defaulted toforkwhich the only value supported on Google Colab. On GCP, other values supported includespawn. Again, see thePytorch XLA pagefor 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 modeIn order to run in single TPU core mode, thefastai_xla_extensionspackage added a methodto_xlato theLearnerclass.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 samefitmethods inLearnerthat 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 ExampleBuild a MNIST classifier -- adapted from fastai courseLesson 4 notebookLoad MNIST datasetpath=untar_data(URLs.MNIST_TINY)Create Fastai DataBlockdatablock=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 transformbatch_tfms=[Normalize.from_stats(*imagenet_stats)])#colabdatablock.summary(path)Setting-up type transforms pipelines
Collecting items from /root/.fastai/data/mnist_tiny
Found 1428 items
2 datasets of sizes 709,699
Setting up Pipeline: PILBase.create
Setting up Pipeline: parent_label -> Categorize -- {'vocab': None, 'sort': True, 'add_na': False}

Building one sample
  Pipeline: PILBase.create
    starting from
      /root/.fastai/data/mnist_tiny/train/7/8791.png
    applying PILBase.create gives
      PILImage mode=RGB size=28x28
  Pipeline: parent_label -> Categorize -- {'vocab': None, 'sort': True, 'add_na': False}
    starting from
      /root/.fastai/data/mnist_tiny/train/7/8791.png
    applying parent_label gives
      7
    applying Categorize -- {'vocab': None, 'sort': True, 'add_na': False} gives
      TensorCategory(1)

Final sample: (PILImage mode=RGB size=28x28, TensorCategory(1))


Collecting items from /root/.fastai/data/mnist_tiny
Found 1428 items
2 datasets of sizes 709,699
Setting up Pipeline: PILBase.create
Setting up Pipeline: parent_label -> Categorize -- {'vocab': None, 'sort': True, 'add_na': False}
Setting up after_item: Pipeline: Resize -- {'size': (28, 28), 'method': 'crop', 'pad_mode': 'reflection', 'resamples': (2, 0), 'p': 1.0} -> ToTensor
Setting up before_batch: Pipeline: 
Setting up after_batch: Pipeline: IntToFloatTensor -- {'div': 255.0, 'div_mask': 1} -> Normalize -- {'mean': tensor([[[[0.4850]],

         [[0.4560]],

         [[0.4060]]]]), 'std': tensor([[[[0.2290]],

         [[0.2240]],

         [[0.2250]]]]), 'axes': (0, 2, 3)}

Building one batch
Applying item_tfms to the first sample:
  Pipeline: Resize -- {'size': (28, 28), 'method': 'crop', 'pad_mode': 'reflection', 'resamples': (2, 0), 'p': 1.0} -> ToTensor
    starting from
      (PILImage mode=RGB size=28x28, TensorCategory(1))
    applying Resize -- {'size': (28, 28), 'method': 'crop', 'pad_mode': 'reflection', 'resamples': (2, 0), 'p': 1.0} gives
      (PILImage mode=RGB size=28x28, TensorCategory(1))
    applying ToTensor gives
      (TensorImage of size 3x28x28, TensorCategory(1))

Adding the next 3 samples

No before_batch transform to apply

Collating items in a batch

Applying batch_tfms to the batch built
  Pipeline: IntToFloatTensor -- {'div': 255.0, 'div_mask': 1} -> Normalize -- {'mean': tensor([[[[0.4850]],

         [[0.4560]],

         [[0.4060]]]]), 'std': tensor([[[[0.2290]],

         [[0.2240]],

         [[0.2250]]]]), 'axes': (0, 2, 3)}
    starting from
      (TensorImage of size 4x3x28x28, TensorCategory([1, 1, 1, 1]))
    applying IntToFloatTensor -- {'div': 255.0, 'div_mask': 1} gives
      (TensorImage of size 4x3x28x28, TensorCategory([1, 1, 1, 1]))
    applying Normalize -- {'mean': tensor([[[[0.4850]],

         [[0.4560]],

         [[0.4060]]]]), 'std': tensor([[[[0.2290]],

         [[0.2240]],

         [[0.2250]]]]), 'axes': (0, 2, 3)} gives
      (TensorImage of size 4x3x28x28, TensorCategory([1, 1, 1, 1]))Create the dataloaderdls=datablock.dataloaders(path,bs=8)# use a small batch size as mnist_tiny dataset is very small#colabdls.show_batch()Create a Fastai CNN Learnerlearner=cnn_learner(dls,resnet18,metrics=accuracy,concat_pool=False)# fastai AdaptivePooling causes a lowering which slows down the training so using nn.AvgPool2DDownloading: "https://download.pytorch.org/models/resnet18-5c106cde.pth" to /root/.cache/torch/hub/checkpoints/resnet18-5c106cde.pth#colablearner.summary()Using thelr_findworks#colablearner.xla_lr_find()start fit









SuggestedLRs(lr_min=7.585775847473997e-08, lr_steep=6.309573450380412e-07)Run one cycle training.#colablearner.xla_fit_one_cycle(5,lr_max=slice(1e-4,0.02))start fitepochtrain_lossvalid_lossaccuracytime00.1900640.3917550.84801100:1710.2732100.2886290.88210200:0620.3221630.4309950.87500000:0630.3544220.5345260.83096600:0640.3723280.5170530.83380700:06Further fine-tuning#colablearner.xla_fit_one_cycle(5,lr_max=slice(7e-4,1e-3))start fitepochtrain_lossvalid_lossaccuracytime00.0467000.4511590.86789800:1410.0564590.3407510.89630700:0720.0618110.3760110.89204500:0630.0712140.3945410.87926100:0640.1375220.4073960.89914800:06Model params are using CPU (since it is the spawned process models that are moved to TPU)#colabone_param(learner.model).devicedevice(type='cpu')Plot loss seems to be working fine.#colablearner.recorder.plot_loss()Single TPU Core ExampleBuild a MNIST classifier -- adapted from fastai courseLesson 4 notebookLoad MNIST datasetpath=untar_data(URLs.MNIST_TINY)Create Fastai DataBlockdatablock=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 TPUbatch_tfms=aug_transforms(do_flip=False,min_scale=0.8))#colabdatablock.summary(path)Setting-up type transforms pipelines
Collecting items from /root/.fastai/data/mnist_tiny
Found 1428 items
2 datasets of sizes 709,699
Setting up Pipeline: PILBase.create
Setting up Pipeline: parent_label -> Categorize -- {'vocab': None, 'sort': True, 'add_na': False}

Building one sample
  Pipeline: PILBase.create
    starting from
      /root/.fastai/data/mnist_tiny/train/7/8791.png
    applying PILBase.create gives
      PILImage mode=RGB size=28x28
  Pipeline: parent_label -> Categorize -- {'vocab': None, 'sort': True, 'add_na': False}
    starting from
      /root/.fastai/data/mnist_tiny/train/7/8791.png
    applying parent_label gives
      7
    applying Categorize -- {'vocab': None, 'sort': True, 'add_na': False} gives
      TensorCategory(1)

Final sample: (PILImage mode=RGB size=28x28, TensorCategory(1))


Collecting items from /root/.fastai/data/mnist_tiny
Found 1428 items
2 datasets of sizes 709,699
Setting up Pipeline: PILBase.create
Setting up Pipeline: parent_label -> Categorize -- {'vocab': None, 'sort': True, 'add_na': False}
Setting up after_item: Pipeline: Resize -- {'size': (28, 28), 'method': 'crop', 'pad_mode': 'reflection', 'resamples': (2, 0), 'p': 1.0} -> ToTensor
Setting up before_batch: Pipeline: 
Setting up after_batch: Pipeline: IntToFloatTensor -- {'div': 255.0, 'div_mask': 1} -> Warp -- {'magnitude': 0.2, 'p': 1.0, 'draw_x': None, 'draw_y': None, 'size': None, 'mode': 'bilinear', 'pad_mode': 'reflection', 'batch': False, 'align_corners': True, 'mode_mask': 'nearest'} -> RandomResizedCropGPU -- {'size': None, 'min_scale': 0.8, 'ratio': (1, 1), 'mode': 'bilinear', 'valid_scale': 1.0, 'p': 1.0} -> Brightness -- {'max_lighting': 0.2, 'p': 1.0, 'draw': None, 'batch': False}

Building one batch
Applying item_tfms to the first sample:
  Pipeline: Resize -- {'size': (28, 28), 'method': 'crop', 'pad_mode': 'reflection', 'resamples': (2, 0), 'p': 1.0} -> ToTensor
    starting from
      (PILImage mode=RGB size=28x28, TensorCategory(1))
    applying Resize -- {'size': (28, 28), 'method': 'crop', 'pad_mode': 'reflection', 'resamples': (2, 0), 'p': 1.0} gives
      (PILImage mode=RGB size=28x28, TensorCategory(1))
    applying ToTensor gives
      (TensorImage of size 3x28x28, TensorCategory(1))

Adding the next 3 samples

No before_batch transform to apply

Collating items in a batch

Applying batch_tfms to the batch built
  Pipeline: IntToFloatTensor -- {'div': 255.0, 'div_mask': 1} -> Warp -- {'magnitude': 0.2, 'p': 1.0, 'draw_x': None, 'draw_y': None, 'size': None, 'mode': 'bilinear', 'pad_mode': 'reflection', 'batch': False, 'align_corners': True, 'mode_mask': 'nearest'} -> RandomResizedCropGPU -- {'size': None, 'min_scale': 0.8, 'ratio': (1, 1), 'mode': 'bilinear', 'valid_scale': 1.0, 'p': 1.0} -> Brightness -- {'max_lighting': 0.2, 'p': 1.0, 'draw': None, 'batch': False}
    starting from
      (TensorImage of size 4x3x28x28, TensorCategory([1, 1, 1, 1]))
    applying IntToFloatTensor -- {'div': 255.0, 'div_mask': 1} gives
      (TensorImage of size 4x3x28x28, TensorCategory([1, 1, 1, 1]))
    applying Warp -- {'magnitude': 0.2, 'p': 1.0, 'draw_x': None, 'draw_y': None, 'size': None, 'mode': 'bilinear', 'pad_mode': 'reflection', 'batch': False, 'align_corners': True, 'mode_mask': 'nearest'} gives
      (TensorImage of size 4x3x28x28, TensorCategory([1, 1, 1, 1]))
    applying RandomResizedCropGPU -- {'size': None, 'min_scale': 0.8, 'ratio': (1, 1), 'mode': 'bilinear', 'valid_scale': 1.0, 'p': 1.0} gives
      (TensorImage of size 4x3x26x26, TensorCategory([1, 1, 1, 1]))
    applying Brightness -- {'max_lighting': 0.2, 'p': 1.0, 'draw': None, 'batch': False} gives
      (TensorImage of size 4x3x26x26, TensorCategory([1, 1, 1, 1]))Create the dataloaderdls=datablock.dataloaders(path)#colabdls.show_batch()Create a Fastai CNN Learnerlearner=cnn_learner(dls,resnet18,metrics=accuracy)#colablearner.summary()Set Learner to XLA mode
This will setup the learner to use the XLA Device#colablearner.to_xla()<fastai.learner.Learner at 0x7fc75a5a8090>Using thelr_findworks#colablearner.lr_find()SuggestedLRs(lr_min=0.02089296132326126, lr_steep=0.0030199517495930195)Run one cycle training.#colablearner.fit_one_cycle(5,lr_max=slice(1e-4,0.02))epochtrain_lossvalid_lossaccuracytime00.8962381.8763740.50500700:5110.7252751.0186290.53218900:1320.6143003.2923560.51359100:0330.5491662.2624870.58083000:0340.5351611.7607730.66094400:03Further fine-tuning#colablearner.fit_one_cycle(5,slice(7e-4,1e-3))epochtrain_lossvalid_lossaccuracytime00.2803590.9924680.68812600:1210.3151321.0135330.71244600:0320.3061870.7452020.71101600:0330.2822780.7012460.73676700:0340.2884740.8543700.71101600:03Model params are using TPU#colabone_param(learner.model).devicedevice(type='xla', index=1)Dataloader device is set to None (as it is the device mover transform that moves the input to the TPU)learner.dls.deviceisNoneTruePlot loss seems to be working fine.#colablearner.recorder.plot_loss()Performance troubleshootingTo check if your model is hiting anaten operation(an operation that is not handled by accelerator device and returned to CPU for default implementation) you can check it with ands then you can report to pytorch xla team.#colabfromfastai_xla_extensions.utilsimportprint_aten_opsprint_aten_ops()needs lowering: Counter: aten::adaptive_max_pool2d
  Value: 311
needs lowering: Counter: aten::adaptive_max_pool2d_backward
  Value: 201SamplesOther examples of fastai notebooks using the fastai_xla_extensions package are also available here:Vision (Multiple TPU Core) - CIFARVision (Multiple TPU Core) - Multi LabelVision (Multiple TPU Core) - PetsVision (Multiple TPU Core) - Plain pytorch training with minimal fastaiVision (Single Core TPU)Collaborative Filtering (Single Core TPU)Tabular (Single Core TPU)Kaggle Notebook version of Vision (Multiple TPU Core) - PetsNEW**Vision (Multi Core) using Pytorch Image Models (timm models) with AlbumentationsMore samples will be added in the future as we fix issues and implement more capabilities.StatusThe 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 usesnbdev-- a jupyter notebook first development environment and is being developed onColab.Development Setupseethis note on how to setup your environment for development