Repository: pytorch/ignite

Issue:
TPU Example / Support
## ❓ Questions/Help/Support

I would like to use Ignite with TPUs and was wondering if there are any examples currently available which currently use TPUs with Ignite or does Ignite have some limitations which would not allow for easy compatibility with TPUs.


Discussion:
@ryanwongsa unfortunately, there is no yet examples of using Ignite with TPUs. 

I think there is no limitations on that as PyTorch xla API is rather similar to GPUs. But, yes, it will be good idea to add some example with Colab link.


Resolving patch:
diff --git a/examples/mnist/mnist_with_tensorboard_on_tpu.py b/examples/mnist/mnist_with_tensorboard_on_tpu.py
new file mode 100644
--- /dev/null
+++ b/examples/mnist/mnist_with_tensorboard_on_tpu.py
@@ -0,0 +1,164 @@
+"""
+ MNIST example with training and validation monitoring using Tensorboard on TPU
+ Requirements:
+    - PyTorch >= 1.5
+    - PyTorch XLA >= 1.5
+    - Tensorboard: `pip install tensorflow` (or just install tensorboard without the rest of tensorflow)
+ Usage:
+    Start tensorboard:
+    ```bash
+    tensorboard --logdir=/tmp/tensorboard_logs/
+    ```
+    Run the example:
+    ```bash
+    python mnist_with_tensorboard_on_tpu.py --log_dir=/tmp/tensorboard_logs
+    ```
+"""
+
+from argparse import ArgumentParser
+
+from torch.utils.data import DataLoader
+from torch import nn
+import torch.nn.functional as F
+from torch.optim import SGD
+from torchvision.datasets import MNIST
+from torchvision.transforms import Compose, ToTensor, Normalize
+
+try:
+    import torch_xla.core.xla_model as xm
+except ImportError:
+    raise RuntimeError(
+        "In order to run PyTorch on TPU we need to install PyTorch XLA:"
+        "\n\t- curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o xla-setup.py"
+        "\n\t- python xla-setup.py --version 1.5"
+    )
+
+from torch.utils.tensorboard import SummaryWriter
+
+from ignite.engine import Events, create_supervised_trainer, create_supervised_evaluator
+from ignite.metrics import Accuracy, Loss, RunningAverage
+
+
+class Net(nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
+        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
+        self.conv2_drop = nn.Dropout2d()
+        self.fc1 = nn.Linear(320, 50)
+        self.fc2 = nn.Linear(50, 10)
+
+    def forward(self, x):
+        x = F.relu(F.max_pool2d(self.conv1(x), 2))
+        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
+        x = x.view(-1, 320)
+        x = F.relu(self.fc1(x))
+        x = F.dropout(x, training=self.training)
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=-1)
+
+
+def get_data_loaders(train_batch_size, val_batch_size):
+    data_transform = Compose([ToTensor(), Normalize((0.1307,), (0.3081,))])
+
+    train_loader = DataLoader(
+        MNIST(download=True, root=".", transform=data_transform, train=True), batch_size=train_batch_size, shuffle=True
+    )
+
+    val_loader = DataLoader(
+        MNIST(download=False, root=".", transform=data_transform, train=False), batch_size=val_batch_size, shuffle=False
+    )
+    return train_loader, val_loader
+
+
+def run(train_batch_size, val_batch_size, epochs, lr, momentum, log_interval, log_dir):
+    train_loader, val_loader = get_data_loaders(train_batch_size, val_batch_size)
+    model = Net()
+    writer = SummaryWriter(log_dir=log_dir)
+
+    # Use TPU device
+    device = xm.xla_device()
+
+    model.to(device)  # Move model before creating optimizer
+    optimizer = SGD(model.parameters(), lr=lr, momentum=momentum)
+
+    # Create trainer and evaluator
+    trainer = create_supervised_trainer(
+        model, optimizer, F.nll_loss, device=device, output_transform=lambda x, y, y_pred, loss: [loss.item(),]
+    )
+
+    evaluator = create_supervised_evaluator(
+        model, metrics={"accuracy": Accuracy(), "nll": Loss(F.nll_loss)}, device=device
+    )
+
+    tracker = xm.RateTracker()
+
+    # Add RateTracker as an output of the training step
+    @trainer.on(Events.ITERATION_COMPLETED)
+    def add_rate_tracker(engine):
+        tracker.add(len(engine.state.batch))
+        engine.state.output.append(tracker.global_rate())
+
+    # Setup output values of the training step as EMA metrics
+    RunningAverage(output_transform=lambda x: x[0]).attach(trainer, "batch_loss")
+    RunningAverage(output_transform=lambda x: x[1]).attach(trainer, "global_rate")
+
+    # Let's log the EMA metrics every `log_interval` iterations
+    @trainer.on(Events.ITERATION_COMPLETED(every=log_interval))
+    def log_training_loss(engine):
+        writer.add_scalar("training/batch_loss", engine.state.metrics["batch_loss"], engine.state.iteration)
+        writer.add_scalar("training/global_rate", engine.state.metrics["global_rate"], engine.state.iteration)
+
+    @trainer.on(Events.EPOCH_COMPLETED)
+    def log_training_results(engine):
+        evaluator.run(train_loader)
+        metrics = evaluator.state.metrics
+        avg_accuracy = metrics["accuracy"]
+        avg_nll = metrics["nll"]
+        print(
+            "Training Results - Epoch: {}  Avg accuracy: {:.2f} Avg loss: {:.2f}".format(
+                engine.state.epoch, avg_accuracy, avg_nll
+            )
+        )
+        writer.add_scalar("training/avg_loss", avg_nll, engine.state.epoch)
+        writer.add_scalar("training/avg_accuracy", avg_accuracy, engine.state.epoch)
+
+    @trainer.on(Events.EPOCH_COMPLETED)
+    def log_validation_results(engine):
+        evaluator.run(val_loader)
+        metrics = evaluator.state.metrics
+        avg_accuracy = metrics["accuracy"]
+        avg_nll = metrics["nll"]
+        print(
+            "Validation Results - Epoch: {}  Avg accuracy: {:.2f} Avg loss: {:.2f}".format(
+                engine.state.epoch, avg_accuracy, avg_nll
+            )
+        )
+        writer.add_scalar("valdation/avg_loss", avg_nll, engine.state.epoch)
+        writer.add_scalar("valdation/avg_accuracy", avg_accuracy, engine.state.epoch)
+
+    # kick everything off
+    trainer.run(train_loader, max_epochs=epochs)
+
+    writer.close()
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--batch_size", type=int, default=64, help="input batch size for training (default: 64)")
+    parser.add_argument(
+        "--val_batch_size", type=int, default=1000, help="input batch size for validation (default: 1000)"
+    )
+    parser.add_argument("--epochs", type=int, default=10, help="number of epochs to train (default: 10)")
+    parser.add_argument("--lr", type=float, default=0.01, help="learning rate (default: 0.01)")
+    parser.add_argument("--momentum", type=float, default=0.5, help="SGD momentum (default: 0.5)")
+    parser.add_argument(
+        "--log_interval", type=int, default=10, help="how many batches to wait before logging training status"
+    )
+    parser.add_argument(
+        "--log_dir", type=str, default="tensorboard_logs", help="log directory for Tensorboard log output"
+    )
+
+    args = parser.parse_args()
+
+    run(args.batch_size, args.val_batch_size, args.epochs, args.lr, args.momentum, args.log_interval, args.log_dir)
diff --git a/ignite/engine/__init__.py b/ignite/engine/__init__.py
--- a/ignite/engine/__init__.py
+++ b/ignite/engine/__init__.py
@@ -48,6 +48,7 @@ def create_supervised_trainer(
         loss_fn (torch.nn loss function): the loss function to use.
         device (str, optional): device type specification (default: None).
             Applies to batches after starting the engine. Model *will not* be moved.
+            Device can be CPU, GPU or TPU.
         non_blocking (bool, optional): if True and this copy is between CPU and GPU, the copy may occur asynchronously
             with respect to the host. For other cases, this argument has no effect.
         prepare_batch (callable, optional): function that receives `batch`, `device`, `non_blocking` and outputs
@@ -74,6 +75,15 @@ def create_supervised_trainer(
         Engine: a trainer engine with supervised update function.
     """
 
+    device_type = device.type if isinstance(device, torch.device) else device
+    on_tpu = "xla" in device_type if device_type is not None else False
+
+    if on_tpu:
+        try:
+            import torch_xla.core.xla_model as xm
+        except ImportError:
+            raise RuntimeError("In order to run on TPU, please install PyTorch XLA")
+
     def _update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[torch.Tensor]]:
         model.train()
         optimizer.zero_grad()
@@ -81,7 +91,12 @@ def _update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[t
         y_pred = model(x)
         loss = loss_fn(y_pred, y)
         loss.backward()
-        optimizer.step()
+
+        if on_tpu:
+            xm.optimizer_step(optimizer, barrier=True)
+        else:
+            optimizer.step()
+
         return output_transform(x, y, y_pred, loss)
 
     trainer = Engine(_update)


Test patch: