Add example for MNIST image classification

Update requirements.
Add notebook for educational purposes.

.gitignore

 __pycache__/
+.ipynb_checkpoints/
+data/
+img/
 *.swp
 *.swo

app/mnist_image_classification.py

+import os
+import numpy as np
+import torch
+import matplotlib.pyplot as plt
+import neural_networks_101.src as src
+
+
+def main():
+    """Main function."""
+    # Random seed for reproducibility
+    torch.manual_seed(42)
+
+    # Get CPU or GPU device for training
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = torch.device(device)
+
+    print(src.misc.time_stamp(), "load MNIST data")
+    train_data, test_data = src.mnist.get_mnist_data()
+
+    # plot some training data
+    print(src.misc.time_stamp(), "plot sample images (random selection)")
+    if not os.path.isdir("../img"):
+        os.makedirs("../img", exist_ok=True)
+    shape = (4, 6)
+    idxs = torch.randint(len(train_data), size=(int(np.prod(shape)),))
+    src.mnist.plot_mnist_data(train_data, idxs, shape=shape)
+    file_name = "../img/mnist_images.png"
+    plt.savefig(file_name, dpi=200)
+    print(src.misc.time_stamp(), f"save to: {file_name}")
+
+    print(src.misc.time_stamp(), "setup NN model")
+    # Send the model to the device (CPU or GPU)
+    model = src.mnist.NeuralNetwork().to(device)
+    # Define the optimizer to user for gradient descent
+    optimizer = torch.optim.Adadelta(model.parameters(), lr=1.0)
+    # Shrinks the learning rate by gamma every step_size
+    scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.7)
+
+    # Train the model
+    n_epochs = 2
+    for epoch in range(n_epochs):
+        print(src.misc.time_stamp(), f"training epoch {epoch+1} / {n_epochs}")
+        with src.misc.timeit("time: {:4.2f} s"):
+            src.mnist.train(model, device, train_data, optimizer,
+                            log_interval=100)
+        test_loss, correct = src.mnist.test(model, device, test_data)
+        rate = 100 * correct / len(test_data)
+        print(src.misc.time_stamp(),
+              f"test set avg. loss: {test_loss}",
+              f" -- accuracy: {correct}/{len(test_data)} ({rate:4.1f} %)")
+        scheduler.step()
+
+    # evaluate forward model
+    print(src.misc.time_stamp(), "evaluate model on training set")
+    # enable evaluation mode, return to training mode with model.train()
+    model.eval()
+    test_loader = torch.utils.data.DataLoader(
+        test_data, batch_size=100, num_workers=1, pin_memory=True,
+        shuffle=True)
+    targets, outputs = [], []
+    with torch.no_grad():  # turn of gradient computation
+        for data, target in test_loader:
+            data, target = data.to(device), target.to(device)
+            output = model(data)
+            targets.append(np.array(target))
+            outputs.append(np.array(output))
+    targets = np.array(targets).reshape(-1, 1)
+    outputs = np.concatenate(outputs, axis=0)
+
+    # plot mislabeled test data
+    print(src.misc.time_stamp(), "plot mislabeled data")
+    # sort indices of test data according to classification error (descending)
+    target_vectors = np.array([np.eye(1, 10, k=int(t)).flatten()
+                               for t in targets])
+    idxs = np.argsort(np.linalg.norm(outputs-target_vectors, axis=1))[::-1]
+    # plot n test samples with largest classification error
+    n = 5
+    src.mnist.plot_mnist_accuracy(
+        test_data, targets, outputs, idxs[:n], figsize=(10, 3*n))
+    file_name = "../img/mnist_mislabeled_data.png"
+    plt.savefig(file_name, dpi=200)
+    print(src.misc.time_stamp(), f"save to: {file_name}")
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

 numpy >= 1.20.0
 scipy >= 1.5.0
-torch >= 1.7.1
 matplotlib >= 3.3.2
+torch >= 1.11.0
+torchvision >= 0.12.0

src/__init__.py

+from . import (
+    misc,
+    target_function,
+    mnist,
+)

src/mnist.py

+"""Utility functions for handling the MNIST data set."""
+from typing import Tuple, List, Optional
+import numpy as np
+import torch
+import torchvision
+import matplotlib.pyplot as plt
+
+
+class NeuralNetwork(torch.nn.Module):
+    """Neural network used for MNIST image classification."""
+
+    def __init__(self) -> None:
+        """Initialize network layers."""
+        super(NeuralNetwork, self).__init__()
+        self.conv1 = torch.nn.Conv2d(
+            1, 32, kernel_size=3, stride=1, padding='valid')
+        self.conv2 = torch.nn.Conv2d(
+            32, 64, kernel_size=3, stride=1, padding='valid')
+        self.dropout1 = torch.nn.Dropout(0.25)
+        self.dropout2 = torch.nn.Dropout(0.5)
+        self.fc1 = torch.nn.Linear(9216, 128)
+        self.fc2 = torch.nn.Linear(128, 10)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Evalute network in input data.
+
+        This function defines the topology of the network.
+
+        Parameters:
+        x : torch.Tensor
+            MNIST input data.
+
+        Returns:
+        :
+            Output vector with probabilities for each class.
+        """
+        x = self.conv1(x)
+        x = torch.nn.functional.relu(x)
+        x = self.conv2(x)
+        x = torch.nn.functional.relu(x)
+        x = torch.nn.functional.max_pool2d(x, 2)
+        x = self.dropout1(x)
+        x = torch.flatten(x, 1)
+        x = self.fc1(x)
+        x = torch.nn.functional.relu(x)
+        x = self.dropout2(x)
+        x = self.fc2(x)
+        output = torch.nn.functional.softmax(x, dim=1)
+        return output
+
+
+def get_mnist_data() -> Tuple[torchvision.datasets.mnist.MNIST,
+                              torchvision.datasets.mnist.MNIST]:
+    """Download the MNIST data set and return data loaders.
+
+    The MNIST data are downloaded to 'data' subdirectory in repository root.
+
+    Returns
+    -------
+    train_loader :
+        DataLoader for training data set.
+    test_loader :
+        DataLoader for test data set.
+    """
+    # The scaled mean and standard deviation of the MNIST dataset
+    # Note: This is precalculated.
+    data_mean = 0.1307
+    data_std = 0.3081
+
+    # Convert input images to tensors and normalize
+    transform = torchvision.transforms.Compose([
+        torchvision.transforms.ToTensor(),
+        torchvision.transforms.Normalize((data_mean,), (data_std,))
+    ])
+
+    # Get the MNIST data from torchvision
+    train_data = torchvision.datasets.MNIST(
+        '../data', train=True, download=True, transform=transform)
+    test_data = torchvision.datasets.MNIST(
+        '../data', train=False, download=True, transform=transform)
+
+    return train_data, test_data
+
+
+def train(model: NeuralNetwork,
+          device: torch.device,
+          train_data: torchvision.datasets.mnist.MNIST,
+          optimizer: torch.optim.Optimizer,
+          log_interval=100
+          ) -> None:
+    """Train the given model.
+
+    Parameters
+    ----------
+    model : NeuralNetwork
+        Neural network model.
+    device : torch.device
+        Hardware to train the model on.
+    train_data : torchvision.datasets.mnist.MNIST
+        Training data the model is trained on.
+    optimizer : torch.optim.Optimizer
+        Optimization employed to train the model.
+    log_interval : int
+        Number of steps the progress is printed after.
+    """
+    kwargs = {'num_workers': 1, 'pin_memory': True, 'shuffle': True}
+    train_loader = torch.utils.data.DataLoader(
+        train_data, batch_size=64, **kwargs)
+    model.train()
+    for batch_idx, (data, target) in enumerate(train_loader):
+        data, target = data.to(device), target.to(device)
+        optimizer.zero_grad()
+        output = model(data)
+        loss = torch.nn.functional.cross_entropy(output, target)
+        loss.backward()
+        optimizer.step()
+        if batch_idx % log_interval == 0:
+            percentage = 100. * batch_idx / len(train_loader)
+            print(f"progress: {percentage:>4.1f} % -- loss: {loss.item()}")
+
+
+def test(model: NeuralNetwork,
+         device: torch.device,
+         test_data: torchvision.datasets.mnist.MNIST
+         ) -> Tuple[float, int]:
+    """Test the network on a test data set.
+
+    Parameters
+    ----------
+    model : NeuralNetwork
+        Neural network model.
+    device : torch.device
+        Hardware to train the model on.
+    test_data : torchvision.datasets.mnist.MNIST
+        Test data set.
+
+    Returns
+    -------
+    test_loss : float
+        Loss on the test data set.
+    correct : int
+        Number of correctly classified test data.
+    """
+    kwargs = {'num_workers': 1, 'pin_memory': True, 'shuffle': True}
+    test_loader = torch.utils.data.DataLoader(
+        test_data, batch_size=16, **kwargs)
+    model.eval()
+    test_loss = 0
+    correct = 0
+    with torch.no_grad():
+        for data, target in test_loader:
+            data, target = data.to(device), target.to(device)
+            output = model(data)
+            # sum up batch loss
+            test_loss += torch.nn.functional.nll_loss(
+                output, target, reduction='sum').item()
+            # get the index of the max log-probability
+            pred = output.argmax(dim=1, keepdim=True)
+            correct += pred.eq(target.view_as(pred)).sum().item()
+    test_loss /= len(test_loader.dataset)
+    return test_loss, correct
+
+
+def plot_mnist_data(data: torchvision.datasets.mnist.MNIST,
+                    idxs: List[int],
+                    shape: Tuple,
+                    **kwargs,
+                    ) -> None:
+    """Plot MNIST images.
+
+    Parameters
+    ----------
+    data : torchvision.datasets.mnist.MNIST
+        MNIST data set.
+    idxs : list of int
+        Indices of data to plot.
+    shape : tuple
+        Specification of arrangement of subplots.
+    """
+    assert np.prod(shape) == len(idxs)
+    layout = [[f"{row}-{col}" for col in range(shape[1])]
+              for row in range(shape[0])]
+    fig, axes = plt.subplot_mosaic(layout, constrained_layout=True, **kwargs)
+    for j, ax in enumerate(axes.values()):
+        img, label = data[idxs[j]]
+        ax.imshow(img.squeeze(), cmap="gray")
+        ax.set_title(f"label = {label}")
+        ax.axis("off")
+
+
+def plot_mnist_accuracy(data: torchvision.datasets.mnist.MNIST,
+                        targets: np.ndarray,
+                        outputs: np.ndarray,
+                        idxs: List[int],
+                        **kwargs
+                        ) -> None:
+    """Plot MNIST images and corresponding classification of the NN model.
+
+    Parameters
+    ----------
+    data : torchvision.datasets.mnist.MNIST
+        MNIST data set.
+    targets : np.ndarray
+        Target values.
+    outputs : np.ndarray
+        Output classifications of model.
+    idxs : list of int
+        Indices of data to plot.
+    """
+    layout = [[f"img-{idx}", f"acc-{idx}"] for idx in idxs]
+    fig, axes = plt.subplot_mosaic(layout, constrained_layout=True, **kwargs)
+    for key, ax in axes.items():
+        idx = int(key[4:])
+        img, label = data[idx]
+        if key[:4] == "img-":
+            ax.imshow(img.squeeze(), cmap="gray")
+            ax.set_title(f"test point = {idx}")
+            ax.axis("off")
+        else:
+            ax.bar(range(10), outputs[idx], log=True, color="gray")
+            ax.axvline(x=label, color="k")
+            ax.set_title(f"target = {label},   "
+                         + f"prediction={np.argmax(outputs[idx])}")