Enable TPU support

.circleci/config.yml

@@ -36,8 +36,10 @@ references:
     run:
       name: Create package
       command: |
+        sudo pip install twine==1.13.0
         python setup.py sdist
         twine check dist/*
+        python setup.py clean
 
   format: &format
    run:

.github/PULL_REQUEST_TEMPLATE.md

@@ -4,6 +4,7 @@
 - [ ] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md)?
 - [ ] Did you make sure to update the docs?   
 - [ ] Did you write any new necessary tests?  
+- [ ] If you made a notable change (that affects users), did you update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CHANGELOG.md)?
 
 ## What does this PR do?
 Fixes # (issue).

.github/workflows/greetings.yml

@@ -10,5 +10,5 @@ jobs:
     - uses: actions/first-interaction@v1
       with:
         repo-token: ${{ secrets.GITHUB_TOKEN }}
-        issue-message: 'Hi''! thanks for your contribution!, great first issue!'
-        pr-message: 'Hey thanks for the input! please give us a bit of time to review it!'
+        issue-message: 'Hey, thanks for your contribution! Great first issue!'
+        pr-message: 'Hey, thanks for the input! Please give us a bit of time to review it!'

docs/source/apex.rst

@@ -1,13 +1,18 @@
 16-bit training
 =================
+Lightning offers 16-bit training for CPUs, GPUs and TPUs.
+
+GPU 16-bit
+-----------
 Lightning uses NVIDIA apex to handle 16-bit precision training.
 
 To use 16-bit precision, do two things:
+
 1. Install Apex
-2. Set the amp trainer flag.
+2. Set the "precision" trainer flag.
 
 Install apex
-----------------------------------------------
+^^^^^^^^^^^^
 .. code-block:: bash
 
     $ git clone https://github.com/NVIDIA/apex
@@ -31,12 +36,25 @@ Install apex
 
 
 Enable 16-bit
---------------
+^^^^^^^^^^^^^
 
 .. code-block:: python
 
-    # DEFAULT
-    trainer = Trainer(amp_level='O1', use_amp=False)
+    # turn on 16-bit
+    trainer = Trainer(amp_level='O1', precision=16)
 
 If you need to configure the apex init for your particular use case or want to use a different way of doing
-16-bit training, override   :meth:`pytorch_lightning.core.LightningModule.configure_apex`.
+16-bit training, override   :meth:`pytorch_lightning.core.LightningModule.configure_apex`.
+
+TPU 16-bit
+----------
+16-bit on TPus is much simpler. To use 16-bit with TPUs set precision to 16 when using the tpu flag
+
+.. code-block:: python
+
+    # DEFAULT
+    trainer = Trainer(num_tpu_cores=8, precision=32)
+
+    # turn on 16-bit
+    trainer = Trainer(num_tpu_cores=8, precision=16)
+

docs/source/experiment_reporting.rst

@@ -38,23 +38,23 @@ To plot metrics into whatever logger you passed in (tensorboard, comet, neptune,
 
 .. code-block:: python
 
-   def training_end(self, batch, batch_idx):
+   def training_end(self, outputs):
       loss = some_loss()
       ...
 
       logs = {'train_loss': loss}
       results = {'log': logs}
       return results
 
-   def validation_end(self, batch, batch_idx):
+   def validation_end(self, outputs):
       loss = some_loss()
       ...
 
       logs = {'val_loss': loss}
       results = {'log': logs}
       return results
 
-   def test_end(self, batch, batch_idx):
+   def test_end(self, outputs):
       loss = some_loss()
       ...
 
@@ -99,7 +99,7 @@ Here we show the validation loss in the progress bar
 
 .. code-block:: python
 
-   def validation_end(self, batch, batch_idx):
+   def validation_end(self, outputs):
       loss = some_loss()
       ...
 

docs/source/index.rst

@@ -55,6 +55,7 @@ PyTorch-Lightning Documentation
    single_gpu
    sequences
    training_tricks
+   tpu
    test_set
    optimizers
    profiler

docs/source/new-project.rst

@@ -60,7 +60,7 @@ Then you could do rapid research by switching between these two and using the sa
     else:
         model = CoolerNotBERT()
 
-    trainer = Trainer(gpus=4, use_amp=True)
+    trainer = Trainer(gpus=4, precision=16)
     trainer.fit(model)
 
 

docs/source/tpu.rst

@@ -0,0 +1,175 @@
+TPU support
+===========
+
+Lightning supports running on TPUs. At this moment, TPUs are only available
+on Google Cloud (GCP). For more information on TPUs
+`watch this video <https://www.youtube.com/watch?v=kPMpmcl_Pyw>`_.
+
+Live demo
+----------
+Check out this `Google Colab <https://colab.research.google.com/drive/1-_LKx4HwAxl5M6xPJmqAAu444LTDQoa3>`_ to see how to train MNIST on TPUs.
+
+TPU Terminology
+---------------
+A TPU is a Tensor processing unit. Each TPU has 8 cores where each
+core is optimized for 128x128 matrix multiplies. In general, a single
+TPU is about as fast as 5 V100 GPUs!
+
+A TPU pod hosts many TPUs on it. Currently, TPU pod v2 has 2048 cores!
+You can request a full pod from Google cloud or a "slice" which gives you
+some subset of those 2048 cores.
+
+How to access TPUs
+-------------------
+To access TPUs there are two main ways.
+
+1. Using google colab.
+2. Using Google Cloud (GCP).
+
+Colab TPUs
+-----------
+Colab is like a jupyter notebook with a free GPU or TPU
+hosted on GCP.
+
+To get a TPU on colab, follow these steps:
+
+1. Go to https://colab.research.google.com/.
+
+2. Click "new notebook" (bottom right of pop-up).
+
+3. Click runtime > change runtime settings. Select Python 3,
+and hardware accelerator "TPU". This will give you a TPU with 8 cores.
+
+4. Next, insert this code into the first cell and execute. This
+will install the xla library that interfaces between PyTorch and
+the TPU.
+
+.. code-block:: python
+
+    import collections
+    from datetime import datetime, timedelta
+    import os
+    import requests
+    import threading
+
+    _VersionConfig = collections.namedtuple('_VersionConfig', 'wheels,server')
+    VERSION = "xrt==1.15.0"  #@param ["xrt==1.15.0", "torch_xla==nightly"]
+    CONFIG = {
+        'xrt==1.15.0': _VersionConfig('1.15', '1.15.0'),
+        'torch_xla==nightly': _VersionConfig('nightly', 'XRT-dev{}'.format(
+            (datetime.today() - timedelta(1)).strftime('%Y%m%d'))),
+    }[VERSION]
+    DIST_BUCKET = 'gs://tpu-pytorch/wheels'
+    TORCH_WHEEL = 'torch-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels)
+    TORCH_XLA_WHEEL = 'torch_xla-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels)
+    TORCHVISION_WHEEL = 'torchvision-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels)
+
+    # Update TPU XRT version
+    def update_server_xrt():
+      print('Updating server-side XRT to {} ...'.format(CONFIG.server))
+      url = 'http://{TPU_ADDRESS}:8475/requestversion/{XRT_VERSION}'.format(
+          TPU_ADDRESS=os.environ['COLAB_TPU_ADDR'].split(':')[0],
+          XRT_VERSION=CONFIG.server,
+      )
+      print('Done updating server-side XRT: {}'.format(requests.post(url)))
+
+    update = threading.Thread(target=update_server_xrt)
+    update.start()
+
+    # Install Colab TPU compat PyTorch/TPU wheels and dependencies
+    !pip uninstall -y torch torchvision
+    !gsutil cp "$DIST_BUCKET/$TORCH_WHEEL" .
+    !gsutil cp "$DIST_BUCKET/$TORCH_XLA_WHEEL" .
+    !gsutil cp "$DIST_BUCKET/$TORCHVISION_WHEEL" .
+    !pip install "$TORCH_WHEEL"
+    !pip install "$TORCH_XLA_WHEEL"
+    !pip install "$TORCHVISION_WHEEL"
+    !sudo apt-get install libomp5
+    update.join()
+5. Once the above is done, install PyTorch Lightning (v 0.6.1+).
+
+.. code-block::
+
+    ! pip install pytorch-lightning
+
+6. Then set up your LightningModule as normal.
+
+7. TPUs require a DistributedSampler. That means you should change your
+train_dataloader (and val, train) code as follows.
+
+.. code-block:: python
+
+    import torch_xla.core.xla_model as xm
+
+    @pl.data_loader
+    def train_dataloader(self):
+        dataset = MNIST(
+            os.getcwd(),
+            train=True,
+            download=True,
+            transform=transforms.ToTensor()
+        )
+
+        # required for TPU support
+        sampler = None
+        if use_tpu:
+            sampler = torch.utils.data.distributed.DistributedSampler(
+                dataset,
+                num_replicas=xm.xrt_world_size(),
+                rank=xm.get_ordinal(),
+                shuffle=True
+            )
+
+        loader = DataLoader(
+            dataset,
+            sampler=sampler,
+            batch_size=32
+        )
+
+        return loader
+
+8. Configure the number of TPU cores in the trainer. You can only choose
+1 or 8. To use a full TPU pod skip to the TPU pod section.
+
+.. code-block:: python
+
+    import pytorch_lightning as pl
+
+    my_model = MyLightningModule()
+    trainer = pl.Trainer(num_tpu_cores=8)
+    trainer.fit(my_model)
+
+That's it! Your model will train on all 8 TPU cores.
+
+TPU Pod
+--------
+To train on more than 8 cores, your code actually doesn't change!
+All you need to do is submit the following command:
+
+.. code-block:: bash
+    $ python -m torch_xla.distributed.xla_dist
+    --tpu=$TPU_POD_NAME
+    --conda-env=torch-xla-nightly
+    -- python /usr/share/torch-xla-0.5/pytorch/xla/test/test_train_imagenet.py --fake_data
+
+16 bit precision
+-----------------
+Lightning also supports training in 16-bit precision with TPUs.
+By default, TPU training will use 32-bit precision. To enable 16-bit, also
+set the 16-bit flag.
+
+.. code-block:: python
+
+    import pytorch_lightning as pl
+
+    my_model = MyLightningModule()
+    trainer = pl.Trainer(num_tpu_cores=8, precision=16)
+    trainer.fit(my_model)
+
+Under the hood the xla library will use the `bfloat16 type <https://en.wikipedia.org/wiki/Bfloat16_floating-point_format>`_.
+
+
+About XLA
+----------
+XLA is the library that interfaces PyTorch with the TPUs.
+For more information check out `XLA <https://github.com/pytorch/xla>`_.

pl_examples/full_examples/semantic_segmentation/models/unet/__init__.py

@@ -0,0 +1,4 @@
+# For relative imports to work in Python 3.6
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.realpath(__file__)))

pl_examples/full_examples/semantic_segmentation/models/unet/model.py

@@ -0,0 +1,42 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from parts import DoubleConv, Down, Up
+
+
+class UNet(nn.Module):
+    '''
+    Architecture based on U-Net: Convolutional Networks for Biomedical Image Segmentation
+    Link - https://arxiv.org/abs/1505.04597
+    '''
+    def __init__(self, num_classes=19, bilinear=False):
+        super().__init__()
+        self.bilinear = bilinear
+        self.num_classes = num_classes
+        self.layer1 = DoubleConv(3, 64)
+        self.layer2 = Down(64, 128)
+        self.layer3 = Down(128, 256)
+        self.layer4 = Down(256, 512)
+        self.layer5 = Down(512, 1024)
+
+        self.layer6 = Up(1024, 512, bilinear=self.bilinear)
+        self.layer7 = Up(512, 256, bilinear=self.bilinear)
+        self.layer8 = Up(256, 128, bilinear=self.bilinear)
+        self.layer9 = Up(128, 64, bilinear=self.bilinear)
+
+        self.layer10 = nn.Conv2d(64, self.num_classes, kernel_size=1)
+
+    def forward(self, x):
+        x1 = self.layer1(x)
+        x2 = self.layer2(x1)
+        x3 = self.layer3(x2)
+        x4 = self.layer4(x3)
+        x5 = self.layer5(x4)
+
+        x6 = self.layer6(x5, x4)
+        x6 = self.layer7(x6, x3)
+        x6 = self.layer8(x6, x2)
+        x6 = self.layer9(x6, x1)
+
+        return self.layer10(x6)