9. Examples

You can find PyTorch examples and tutorials in the Graphcore GitHub examples repository. This contains:

• Examples of popular machine learning models for training and inference

• Tutorials

• Examples of PopTorch and IPU features

• Examples of simple models

• Source code from videos, blogs and other documents

9.1. MNIST example

The example in Listing 9.1 shows how an MNIST model can be run on the IPU. The highlighted lines show the PopTorch-specific code required to run the example on multiple IPUs.

You can download the full source code from GitHub: mnist.py.

To run this example you will need to install the Poplar SDK (see the Getting Started Guide for your IPU system) and the appropriate version of torchvision:

$ python3 -m pip install torchvision==0.11.1

Listing 9.1 MNIST example

import torch
import torch.nn as nn
import torchvision
import poptorch

# Normal pytorch batch size
training_batch_size = 20
validation_batch_size = 100

opts = poptorch.Options()
# Device "step"
opts.deviceIterations(20)

# How many IPUs to replicate over.
opts.replicationFactor(4)

opts.randomSeed(42)

# Load MNIST normally.
training_data = poptorch.DataLoader(
    opts,
    torchvision.datasets.MNIST('mnist_data/',
                               train=True,
                               download=True,
                               transform=torchvision.transforms.Compose([
                                   torchvision.transforms.ToTensor(),
                                   torchvision.transforms.Normalize(
                                       (0.1307, ), (0.3081, ))
                               ])),
    batch_size=training_batch_size,
    shuffle=True)

# Load MNIST normally.
val_options = poptorch.Options()
validation_data = poptorch.DataLoader(
    val_options,
    torchvision.datasets.MNIST('mnist_data/',
                               train=True,
                               download=True,
                               transform=torchvision.transforms.Compose([
                                   torchvision.transforms.ToTensor(),
                                   torchvision.transforms.Normalize(
                                       (0.1307, ), (0.3081, ))
                               ])),
    batch_size=validation_batch_size,
    shuffle=True,
    drop_last=True)

# A helper block to build convolution-pool-relu blocks.
class Block(nn.Module):
    def __init__(self, in_channels, num_filters, kernel_size, pool_size):
        super(Block, self).__init__()
        self.conv = nn.Conv2d(in_channels,
                              num_filters,
                              kernel_size=kernel_size)
        self.pool = nn.MaxPool2d(kernel_size=pool_size)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.conv(x)
        x = self.pool(x)
        x = self.relu(x)
        return x

# Define the network using the above blocks.
class Network(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = Block(1, 10, 5, 2)
        self.layer2 = Block(10, 20, 5, 2)
        self.layer3 = nn.Linear(320, 256)
        self.layer3_act = nn.ReLU()
        self.layer4 = nn.Linear(256, 10)

        self.softmax = nn.LogSoftmax(1)
        self.loss = nn.NLLLoss(reduction="mean")

    def forward(self, x, target=None):
        x = self.layer1(x)
        x = self.layer2(x)
        x = x.view(-1, 320)

        x = self.layer3_act(self.layer3(x))
        x = self.layer4(x)
        x = self.softmax(x)

        if target is not None:
            loss = self.loss(x, target)
            return x, loss
        return x

# Create our model.
model = Network()

# Create model for training which will run on IPU.
training_model = poptorch.trainingModel(model, training_data.options)

# Same model as above, they will share weights (in 'model') which once training is finished can be copied back.
inference_model = poptorch.inferenceModel(model, validation_data.options)

def train():
    for batch_number, (data, labels) in enumerate(training_data):
        output, losses = training_model(data, labels)

        if batch_number % 10 == 0:
            print(f"PoptorchIPU loss at batch: {batch_number} is {losses}")

            # Pick the highest probability.
            _, ind = torch.max(output, 1)
            assert training_model.options.output_mode in (
                poptorch.OutputMode.All, poptorch.OutputMode.Final
            ), "Only 'Final' and 'All' OutputMode supported"
            # If we're using Final: only keep the last labels, no-op if using All
            num_labels = ind.shape[0]
            labels = labels[-num_labels:]
            eq = torch.eq(ind, labels)
            elms, counts = torch.unique(eq,
                                        sorted=False,
                                        return_counts=True)

            acc = 0.0
            if len(elms) == 2:
                if elms[0]:
                    acc = (counts[0].item() / num_labels) * 100.0
                else:
                    acc = (counts[1].item() / num_labels) * 100.0

            print(
                f"Training accuracy: {acc}% from batch of size {num_labels}"
            )
    print("Done training")

def test():
    correct = 0
    total = 0
    with torch.no_grad():
        for (data, labels) in validation_data:
            output = inference_model(data)

            # Argmax the probabilities to get the highest.
            _, ind = torch.max(output, 1)

            # Compare it against the ground truth for this batch.
            eq = torch.eq(ind, labels)

            # Count the number which are True and the number which are False.
            elms, counts = torch.unique(eq,
                                        sorted=False,
                                        return_counts=True)

            if len(elms) == 2 or elms[0]:
                if elms[0]:
                    correct += counts[0].item()
                else:
                    correct += counts[1].item()

            total += validation_batch_size
    print("Validation: of " + str(total) + " samples we got: " +
          str((correct / total) * 100.0) + "% correct")

# Train on IPU.
train()

test()