update mixed precision with comments

Author: aladdinpersson

ML/Pytorch/Basics/pytorch_mixed_precision_example.py

@@ -1,9 +1,23 @@
+"""
+Example code of how to use mixed precision training with PyTorch. In this
+case with a (very) small and simple CNN training on MNIST dataset. This
+example is based on the official PyTorch documentation on mixed precision 
+training.
+
+Programmed by Aladdin Persson <aladdin.persson at hotmail dot com>
+*    2020-04-10 Initial programming
+*    2022-12-19 Updated comments, made sure it works with latest PyTorch
+
+"""
+
 # Imports
 import torch
 import torch.nn as nn  # All neural network modules, nn.Linear, nn.Conv2d, BatchNorm, Loss functions
 import torch.optim as optim  # For all Optimization algorithms, SGD, Adam, etc.
 import torch.nn.functional as F  # All functions that don't have any parameters
-from torch.utils.data import DataLoader  # Gives easier dataset managment and creates mini batches
+from torch.utils.data import (
+    DataLoader,
+)  # Gives easier dataset managment and creates mini batches
 import torchvision.datasets as datasets  # Has standard datasets we can import in a nice way
 import torchvision.transforms as transforms  # Transformations we can perform on our dataset
 
@@ -12,9 +26,21 @@
 class CNN(nn.Module):
     def __init__(self, in_channels=1, num_classes=10):
         super(CNN, self).__init__()
-        self.conv1 = nn.Conv2d(in_channels=1, out_channels=420, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.conv1 = nn.Conv2d(
+            in_channels=1,
+            out_channels=420,
+            kernel_size=(3, 3),
+            stride=(1, 1),
+            padding=(1, 1),
+        )
         self.pool = nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2))
-        self.conv2 = nn.Conv2d(in_channels=420, out_channels=1000, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.conv2 = nn.Conv2d(
+            in_channels=420,
+            out_channels=1000,
+            kernel_size=(3, 3),
+            stride=(1, 1),
+            padding=(1, 1),
+        )
         self.fc1 = nn.Linear(1000 * 7 * 7, num_classes)
 
     def forward(self, x):
@@ -29,7 +55,8 @@ def forward(self, x):
 
 
 # Set device
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+assert device == "cuda", "GPU not available"
 
 # Hyperparameters
 in_channel = 1
@@ -39,9 +66,13 @@ def forward(self, x):
 num_epochs = 5
 
 # Load Data
-train_dataset = datasets.MNIST(root='dataset/', train=True, transform=transforms.ToTensor(), download=True)
+train_dataset = datasets.MNIST(
+    root="dataset/", train=True, transform=transforms.ToTensor(), download=True
+)
 train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
-test_dataset = datasets.MNIST(root='dataset/', train=False, transform=transforms.ToTensor(), download=True)
+test_dataset = datasets.MNIST(
+    root="dataset/", train=False, transform=transforms.ToTensor(), download=True
+)
 test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=True)
 
 # Initialize network
@@ -89,10 +120,12 @@ def check_accuracy(loader, model):
             num_correct += (predictions == y).sum()
             num_samples += predictions.size(0)
 
-        print(f'Got {num_correct} / {num_samples} with accuracy {float(num_correct) / float(num_samples) * 100:.2f}')
+        print(
+            f"Got {num_correct} / {num_samples} with accuracy {float(num_correct) / float(num_samples) * 100:.2f}"
+        )
 
     model.train()
 
 
 check_accuracy(train_loader, model)
-check_accuracy(test_loader, model)
+check_accuracy(test_loader, model)