Issue #64 - Pytorch Tutorial 3 - Jupter Notebook and Markdown

tutorials/pytorch/tut3_mixed_precision/walkthrough.py

@@ -1,30 +1,31 @@
 #!/usr/bin/env python3
-# Copyright (c) 2021 Graphcore Ltd. All rights reserved.
-
+"""
+Copyright (c) 2021 Graphcore Ltd. All rights reserved.
+"""
+"""
 # Half and mixed precision in PopTorch
-
-# This tutorial shows how to use half and mixed precision in PopTorch with the
-# example task of training a simple CNN model on a single
-# Graphcore IPU (Mk1 or Mk2).
-
-# Requirements:
-#   - an installed Poplar SDK. See the Getting Started guide for your IPU
-#     hardware for details of how to install the SDK;
-#   - Other Python modules: `pip install -r requirements.txt`
-
-# Import the packages
+This tutorial shows how to use half and mixed precision in PopTorch with the
+example task of training a simple CNN model on a single
+Graphcore IPU (Mk1 or Mk2).
+"""
+"""
+Requirements:
+- an installed Poplar SDK. See the Getting Started guide for your IPU
+hardware for details of how to install the SDK;
+- Other Python modules: `pip install -r requirements.txt`
+"""
+"""
+Import the packages
+"""
 import torch
 import torch.nn as nn
-
 import torchvision
 import torchvision.transforms as transforms
-
 import poptorch
-import argparse
-from tqdm import tqdm
-
-
-# Build the model
+from tqdm.auto import tqdm
+"""
+Build the model
+"""
 class CustomModel(nn.Module):
     def __init__(self):
         super().__init__()
@@ -50,22 +51,34 @@ def forward(self, x, labels=None):
             return x, self.loss(x, labels)
         return x
 
-model = CustomModel()
+"""
+Choose parameters
+"""
+# Cast the model parameters to FP16
+model_half = True
 
-parser = argparse.ArgumentParser()
-parser.add_argument('--model-half', dest='model_half', action='store_true', help='Cast the model parameters to FP16')
-parser.add_argument('--data-half', dest='data_half', action='store_true', help='Cast the data to FP16')
-parser.add_argument('--optimizer-half', dest='optimizer_half', action='store_true', help='Cast the accumulation type of the optimiser to FP16')
-parser.add_argument('--stochastic-rounding', dest='stochastic_rounding', action='store_true', help='Use stochasting rounding')
-parser.add_argument('--partials-half', dest='partials_half', action='store_true', help='Set partials data type to FP16')
-args = parser.parse_args()
+# Cast the data to FP16
+data_half = True
 
-# Casting a model's parameters
-if args.model_half:
-    model = model.half()
+# Cast the accumulation type of the optimiser to FP16
+optimizer_half = True
+
+# Use stochasting rounding
+stochastic_rounding = True
 
-# Prepare the data
-if args.data_half:
+# Set partials data type to FP16
+partials_half = True
+"""
+Casting a model's parameters
+"""
+model = CustomModel()
+
+if model_half:
+    model = model.half()
+"""
+Prepare the data
+"""
+if data_half:
     transform = transforms.Compose([transforms.Resize(128),
                                     transforms.ToTensor(),
                                     transforms.Normalize((0.5,), (0.5,)),
@@ -83,9 +96,10 @@ def forward(self, x, labels=None):
                                                  transform=transform,
                                                  download=True,
                                                  train=False)
-
-# Optimizer and loss scaling
-if args.optimizer_half:
+"""
+Optimizer and loss scaling
+"""
+if optimizer_half:
     optimizer = poptorch.optim.AdamW(model.parameters(),
                                      lr=0.001,
                                      loss_scaling=1024,
@@ -94,50 +108,67 @@ def forward(self, x, labels=None):
     optimizer = poptorch.optim.AdamW(model.parameters(),
                                      lr=0.001,
                                      accum_type=torch.float32)
-
-
-# Set PopTorch's options
+"""
+Set PopTorch's options
+"""
 opts = poptorch.Options()
-
-# Stochastic rounding
-if args.stochastic_rounding:
+"""
+Stochastic rounding
+"""
+if stochastic_rounding:
     opts.Precision.enableStochasticRounding(True)
-# Partials data type
-if args.partials_half:
+"""
+Partials data type
+"""
+if partials_half:
     opts.Precision.setPartialsType(torch.half)
 else:
     opts.Precision.setPartialsType(torch.float)
-
-# Train the model
+"""
+Initiate DataLoader and the model training PopTorch wrapper
+"""
 train_dataloader = poptorch.DataLoader(opts,
                                        train_dataset,
                                        batch_size=12,
                                        shuffle=True,
                                        num_workers=40)
+
 poptorch_model = poptorch.trainingModel(model,
                                         options=opts,
                                         optimizer=optimizer)
-
+"""
+Execute training
+"""
 epochs = 10
 for epoch in tqdm(range(epochs), desc="epochs"):
     total_loss = 0.0
     for data, labels in tqdm(train_dataloader, desc="batches", leave=False):
         output, loss = poptorch_model(data, labels)
         total_loss += loss
-
-# Evaluate the model
+poptorch_model.detachFromDevice()
+# sst_hide_output
+"""
+Evaluate the model - initiate the inference PopTorch wrapper and a new DataLoader for the test data with labels
+"""
 model.eval()
 poptorch_model_inf = poptorch.inferenceModel(model, options=opts)
 test_dataloader = poptorch.DataLoader(opts,
                                       test_dataset,
                                       batch_size=32,
                                       num_workers=40)
-
+"""
+Run inference on the labelled data
+"""
 predictions, labels = [], []
 for data, label in test_dataloader:
     predictions += poptorch_model_inf(data).data.float().max(dim=1).indices
     labels += label
+poptorch_model_inf.detachFromDevice()
+# sst_hide_output
 
+"""
+Show accuracy
+"""
 print(f"""Eval accuracy on IPU: {100 *
                 (1 - torch.count_nonzero(torch.sub(torch.tensor(labels),
                 torch.tensor(predictions))) / len(labels)):.2f}%""")