Metadata-Version: 2.1
Name: fastai_xla_extensions
Version: 0.0.5
Summary: A library to extend fastai to run on TPUs using pytorch-xla
Home-page: https://github.com/butchland/fastai_xla_extensions/tree/master/
Author: Butch Landingin
Author-email: butchland@gmail.com
License: Apache Software License 2.0
Description: # FastAI XLA Extensions Library
        > The FastAI XLA Extensions library package allows your fastai/Pytorch models to run on TPUs using the Pytorch-XLA library.
        
        
        ## Install
        
        `pip install git+https://github.com/butchland/fastai_xla_extensions`
        
        ## How to use
        
        ### 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](https://pytorch.org/xla/release/1.7/index.html).
        
        If running on Colab, make sure the Runtime Type is set to TPU.
        
        
        ## Install fastai
        
        Use the latest fastai and fastcore versions
        
        ```
        #hide_output
        #colab
        !pip install -Uqq fastcore --upgrade
        !pip install -Uqq fastai --upgrade
        ```
        
        ## Install Pytorch XLA
        
        
        This is the official way to install Pytorch-XLA 1.7 as per the [instructions here](https://colab.research.google.com/github/pytorch/xla/blob/master/contrib/colab/getting-started.ipynb#scrollTo=CHzziBW5AoZH)
        
        ```
        #hide_output
        #colab
        !pip install -Uqq cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.7-cp36-cp36m-linux_x86_64.whl
        ```
        
        ## Check if XLA is available
        
        ### Import the libraries
        Import the fastai and fastai_xla_extensions libraries
        
        ```
        #colab
        #hide_output
        import fastai_xla_extensions.core
        ```
        
        ```
        from fastai.vision.all import *
        ```
        
        ### Example
        Build a MNIST classifier -- adapted from fastai course [Lesson 4 notebook](https://github.com/fastai/course-v4/blob/master/nbs/04_mnist_basics.ipynb)
        
        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_tfms=aug_transforms(do_flip=False,min_scale=0.8)
        )
        ```
        
        ```
        #colab
        datablock.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/7770.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/7770.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 4x3x27x27, TensorCategory([1, 1, 1, 1]))
                applying Brightness -- {'max_lighting': 0.2, 'p': 1.0, 'draw': None, 'batch': False} gives
                  (TensorImage of size 4x3x27x27, TensorCategory([1, 1, 1, 1]))
        
        
        Create the dataloader
        
        ```
        dls = datablock.dataloaders(path)
        ```
        
        ```
        #colab
        dls.show_batch()
        ```
        
        
        ![png](docs/images/output_21_0.png)
        
        
        Create a Fastai CNN Learner
        
        
        ```
        learner = cnn_learner(dls, resnet18, metrics=accuracy)
                              
        ```
        
        ```
        #colab
        learner.summary()
        ```
        
        
        <table border="1" class="dataframe">
          <thead>
            <tr style="text-align: left;">
              <th>epoch</th>
              <th>train_loss</th>
              <th>valid_loss</th>
              <th>accuracy</th>
              <th>time</th>
            </tr>
          </thead>
          <tbody>
            <tr>
              <td>0</td>
              <td>None</td>
              <td>None</td>
              <td>00:00</td>
            </tr>
          </tbody>
        </table>
        
        
        
        
        
            Sequential (Input shape: ['64 x 3 x 28 x 28'])
            ================================================================
            Layer (type)         Output Shape         Param #    Trainable 
            ================================================================
            Conv2d               64 x 64 x 14 x 14    9,408      True      
            ________________________________________________________________
            BatchNorm2d          64 x 64 x 14 x 14    128        True      
            ________________________________________________________________
            ReLU                 64 x 64 x 14 x 14    0          False     
            ________________________________________________________________
            MaxPool2d            64 x 64 x 7 x 7      0          False     
            ________________________________________________________________
            Conv2d               64 x 64 x 7 x 7      36,864     True      
            ________________________________________________________________
            BatchNorm2d          64 x 64 x 7 x 7      128        True      
            ________________________________________________________________
            ReLU                 64 x 64 x 7 x 7      0          False     
            ________________________________________________________________
            Conv2d               64 x 64 x 7 x 7      36,864     True      
            ________________________________________________________________
            BatchNorm2d          64 x 64 x 7 x 7      128        True      
            ________________________________________________________________
            Conv2d               64 x 64 x 7 x 7      36,864     True      
            ________________________________________________________________
            BatchNorm2d          64 x 64 x 7 x 7      128        True      
            ________________________________________________________________
            ReLU                 64 x 64 x 7 x 7      0          False     
            ________________________________________________________________
            Conv2d               64 x 64 x 7 x 7      36,864     True      
            ________________________________________________________________
            BatchNorm2d          64 x 64 x 7 x 7      128        True      
            ________________________________________________________________
            Conv2d               64 x 128 x 4 x 4     73,728     True      
            ________________________________________________________________
            BatchNorm2d          64 x 128 x 4 x 4     256        True      
            ________________________________________________________________
            ReLU                 64 x 128 x 4 x 4     0          False     
            ________________________________________________________________
            Conv2d               64 x 128 x 4 x 4     147,456    True      
            ________________________________________________________________
            BatchNorm2d          64 x 128 x 4 x 4     256        True      
            ________________________________________________________________
            Conv2d               64 x 128 x 4 x 4     8,192      True      
            ________________________________________________________________
            BatchNorm2d          64 x 128 x 4 x 4     256        True      
            ________________________________________________________________
            Conv2d               64 x 128 x 4 x 4     147,456    True      
            ________________________________________________________________
            BatchNorm2d          64 x 128 x 4 x 4     256        True      
            ________________________________________________________________
            ReLU                 64 x 128 x 4 x 4     0          False     
            ________________________________________________________________
            Conv2d               64 x 128 x 4 x 4     147,456    True      
            ________________________________________________________________
            BatchNorm2d          64 x 128 x 4 x 4     256        True      
            ________________________________________________________________
            Conv2d               64 x 256 x 2 x 2     294,912    True      
            ________________________________________________________________
            BatchNorm2d          64 x 256 x 2 x 2     512        True      
            ________________________________________________________________
            ReLU                 64 x 256 x 2 x 2     0          False     
            ________________________________________________________________
            Conv2d               64 x 256 x 2 x 2     589,824    True      
            ________________________________________________________________
            BatchNorm2d          64 x 256 x 2 x 2     512        True      
            ________________________________________________________________
            Conv2d               64 x 256 x 2 x 2     32,768     True      
            ________________________________________________________________
            BatchNorm2d          64 x 256 x 2 x 2     512        True      
            ________________________________________________________________
            Conv2d               64 x 256 x 2 x 2     589,824    True      
            ________________________________________________________________
            BatchNorm2d          64 x 256 x 2 x 2     512        True      
            ________________________________________________________________
            ReLU                 64 x 256 x 2 x 2     0          False     
            ________________________________________________________________
            Conv2d               64 x 256 x 2 x 2     589,824    True      
            ________________________________________________________________
            BatchNorm2d          64 x 256 x 2 x 2     512        True      
            ________________________________________________________________
            Conv2d               64 x 512 x 1 x 1     1,179,648  True      
            ________________________________________________________________
            BatchNorm2d          64 x 512 x 1 x 1     1,024      True      
            ________________________________________________________________
            ReLU                 64 x 512 x 1 x 1     0          False     
            ________________________________________________________________
            Conv2d               64 x 512 x 1 x 1     2,359,296  True      
            ________________________________________________________________
            BatchNorm2d          64 x 512 x 1 x 1     1,024      True      
            ________________________________________________________________
            Conv2d               64 x 512 x 1 x 1     131,072    True      
            ________________________________________________________________
            BatchNorm2d          64 x 512 x 1 x 1     1,024      True      
            ________________________________________________________________
            Conv2d               64 x 512 x 1 x 1     2,359,296  True      
            ________________________________________________________________
            BatchNorm2d          64 x 512 x 1 x 1     1,024      True      
            ________________________________________________________________
            ReLU                 64 x 512 x 1 x 1     0          False     
            ________________________________________________________________
            Conv2d               64 x 512 x 1 x 1     2,359,296  True      
            ________________________________________________________________
            BatchNorm2d          64 x 512 x 1 x 1     1,024      True      
            ________________________________________________________________
            AdaptiveAvgPool2d    64 x 512 x 1 x 1     0          False     
            ________________________________________________________________
            AdaptiveMaxPool2d    64 x 512 x 1 x 1     0          False     
            ________________________________________________________________
            Flatten              64 x 1024            0          False     
            ________________________________________________________________
            BatchNorm1d          64 x 1024            2,048      True      
            ________________________________________________________________
            Dropout              64 x 1024            0          False     
            ________________________________________________________________
            Linear               64 x 512             524,288    True      
            ________________________________________________________________
            ReLU                 64 x 512             0          False     
            ________________________________________________________________
            BatchNorm1d          64 x 512             1,024      True      
            ________________________________________________________________
            Dropout              64 x 512             0          False     
            ________________________________________________________________
            Linear               64 x 2               1,024      True      
            ________________________________________________________________
            
            Total params: 11,704,896
            Total trainable params: 11,704,896
            Total non-trainable params: 0
            
            Optimizer used: <function Adam at 0x7f1399bb1730>
            Loss function: FlattenedLoss of CrossEntropyLoss()
            
            Callbacks:
              - TrainEvalCallback
              - XLAOptCallback
              - Recorder
              - ProgressCallback
        
        
        
        ### Set Learner to XLA mode
        This will setup the learner to use the XLA Device
        
        ```
        #colab
        learner.to_xla()
        ```
        
        
        
        
            <fastai.learner.Learner at 0x7fb6809e7f28>
        
        
        
        Using the `lr_find` works 
        
        ```
        #colab
        learner.lr_find()
        ```
        
        
        
        
        
        
        
        
            SuggestedLRs(lr_min=0.02089296132326126, lr_steep=0.0010000000474974513)
        
        
        
        
        ![png](docs/images/output_28_2.png)
        
        
        Run one cycle training.
        
        
        ```
        #colab
        learner.fit_one_cycle(5,lr_max=slice(1e-4,0.02))
        ```
        
        
        <table border="1" class="dataframe">
          <thead>
            <tr style="text-align: left;">
              <th>epoch</th>
              <th>train_loss</th>
              <th>valid_loss</th>
              <th>accuracy</th>
              <th>time</th>
            </tr>
          </thead>
          <tbody>
            <tr>
              <td>0</td>
              <td>0.749318</td>
              <td>0.348496</td>
              <td>0.861230</td>
              <td>00:12</td>
            </tr>
            <tr>
              <td>1</td>
              <td>0.634834</td>
              <td>0.676399</td>
              <td>0.791130</td>
              <td>00:03</td>
            </tr>
            <tr>
              <td>2</td>
              <td>0.621724</td>
              <td>0.506193</td>
              <td>0.834049</td>
              <td>00:03</td>
            </tr>
            <tr>
              <td>3</td>
              <td>0.553649</td>
              <td>0.503763</td>
              <td>0.824034</td>
              <td>00:03</td>
            </tr>
            <tr>
              <td>4</td>
              <td>0.500875</td>
              <td>0.435466</td>
              <td>0.846924</td>
              <td>00:03</td>
            </tr>
          </tbody>
        </table>
        
        
        Further fine-tuning
        
        ```
        #colab
        learner.fit_one_cycle(5,slice(7e-4, 1e-3))
        ```
        
        
        <table border="1" class="dataframe">
          <thead>
            <tr style="text-align: left;">
              <th>epoch</th>
              <th>train_loss</th>
              <th>valid_loss</th>
              <th>accuracy</th>
              <th>time</th>
            </tr>
          </thead>
          <tbody>
            <tr>
              <td>0</td>
              <td>0.355362</td>
              <td>0.312893</td>
              <td>0.889843</td>
              <td>00:07</td>
            </tr>
            <tr>
              <td>1</td>
              <td>0.361367</td>
              <td>0.271736</td>
              <td>0.895565</td>
              <td>00:03</td>
            </tr>
            <tr>
              <td>2</td>
              <td>0.324752</td>
              <td>0.205274</td>
              <td>0.932761</td>
              <td>00:03</td>
            </tr>
            <tr>
              <td>3</td>
              <td>0.282965</td>
              <td>0.197236</td>
              <td>0.925608</td>
              <td>00:03</td>
            </tr>
            <tr>
              <td>4</td>
              <td>0.246091</td>
              <td>0.187480</td>
              <td>0.939914</td>
              <td>00:03</td>
            </tr>
          </tbody>
        </table>
        
        
        Model params are using TPU
        
        ```
        #colab
        one_param(learner.model).device
        ```
        
        
        
        
            device(type='xla', index=1)
        
        
        
        Plot loss seems to be working fine.
        
        ```
        #colab
        learner.recorder.plot_loss()
        ```
        
        
        ![png](docs/images/output_36_0.png)
        
        
        ## Samples
        
        Other examples of fastai notebooks using the fastai_xla_extensions package are also available here: 
        
        * [Vision](https://github.com/butchland/fastai_xla_extensions/blob/master/samples/MNIST_TPU_demo.ipynb) [![](https://camo.githubusercontent.com/52feade06f2fecbf006889a904d221e6a730c194/68747470733a2f2f636f6c61622e72657365617263682e676f6f676c652e636f6d2f6173736574732f636f6c61622d62616467652e737667)](https://colab.research.google.com/github/butchland/fastai_xla_extensions/blob/master/samples/MNIST_TPU_demo.ipynb)
        
        * [Collaborative Filtering](https://github.com/butchland/fastai_xla_extensions/blob/master/samples/MovieLensCollaborativeFilteringTPU.ipynb) [![](https://camo.githubusercontent.com/52feade06f2fecbf006889a904d221e6a730c194/68747470733a2f2f636f6c61622e72657365617263682e676f6f676c652e636f6d2f6173736574732f636f6c61622d62616467652e737667)](https://colab.research.google.com/github/butchland/fastai_xla_extensions/blob/master/samples/MovieLensCollaborativeFilteringTPU.ipynb)
        
        * [Tabular](https://github.com/butchland/fastai_xla_extensions/blob/master/samples/TabularTrainingTPUdemo.ipynb) [![](https://camo.githubusercontent.com/52feade06f2fecbf006889a904d221e6a730c194/68747470733a2f2f636f6c61622e72657365617263682e676f6f676c652e636f6d2f6173736574732f636f6c61622d62616467652e737667)](https://colab.research.google.com/github/butchland/fastai_xla_extensions/blob/master/samples/TabularTrainingTPUdemo.ipynb)
        
        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 very early development phase (not even alpha) which means there's still a lot of things not working. 
        
        Use it at your own risk.
        
        If you wish to contribute to the project, fork it and make pull request. 
        
        This project uses [nbdev](https://nbdev.fast.ai/) -- a jupyter notebook first development environment and is being developed on [Colab](https://colab.research.google.com).
        
        
Keywords: fastai pytorch pytorch-xla tpu
Platform: UNKNOWN
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Natural Language :: English
Classifier: Programming Language :: Python :: 3.6
Classifier: Programming Language :: Python :: 3.7
Classifier: Programming Language :: Python :: 3.8
Requires-Python: >=3.6
Description-Content-Type: text/markdown
