src.train

  1import matplotlib.pyplot as plt
  2import numpy as np
  3import torch
  4from rich import print
  5from rich.progress import track
  6from torch import nn
  7from torch.nn import MSELoss
  8from torch.optim import Optimizer
  9from torch.optim import SGD
 10from torch.utils.data import DataLoader
 11from torch.utils.tensorboard.writer import SummaryWriter
 12
 13
 14def training_iteration(
 15    model: nn.Module,
 16    loss_fn: nn.Module,
 17    optimizer: Optimizer,
 18    dataloader_train: DataLoader,
 19    global_step: int,
 20    writer: SummaryWriter,
 21) -> np.floating:
 22    """Perform a single training iteration on the DataLoader.  Average losses are recorded and saved to Tensorboard logs.
 23
 24    Args:
 25        model (nn.Module): Torch model.
 26        loss_fn (nn.Module): Torch loss.
 27        optimizer (Optimizer): Torch optimizer.
 28        dataloader_train (DataLoader): DataLoader that outputs `(x, y)` where `x` is input and `y` is the true label.
 29        global_step (int): The global step for tracking this training iteration.
 30        writer (SummaryWriter): Tensorboard writer.
 31
 32    Returns:
 33        np.floating: Average training loss for this iteration.
 34    """
 35    # Initialization
 36    train_losses = []
 37
 38    for x, y in track(dataloader_train, description="training"):
 39        # fmt: off
 40        y_hat, _ = model(x)                 # predict
 41        loss = loss_fn(y_hat, y.long())     # calculate loss
 42        optimizer.zero_grad()               # zero the gradients
 43        loss.backward()                     # backpropagation
 44        optimizer.step()                    # update the weights
 45        train_losses.append(loss.item())    # save off the loss
 46        # fmt: on
 47
 48    # Record the average training loss for this training iteration
 49    avg_train_loss = np.average(train_losses)
 50    writer.add_scalar(tag="Loss/train", scalar_value=avg_train_loss, global_step=global_step)
 51
 52    return avg_train_loss
 53
 54
 55def generate_canocial(
 56    model: nn.Module,
 57    model_weights_path: str,
 58    n_epochs: int = 10,
 59    n_steps: int = 100,
 60):
 61    r"""The idea behind a canocial input is that it should be the input that maximizes the predicted value of a network.  Think of this as `inverse` training rather than standard training.  You want to specify an output and ask the network what is the corresponding input.
 62
 63    The way to achieve this is by freezing all of the layers of a network and setting $y_{true}$ to your target class.  When performing backpropagation via $loss=\hat{y} - y_{true}$ you backpropagate the gradients all the way to the input.
 64
 65    The theory is that once the losses converge to zero, the input should represent the maximum values needed to push the predictions of your neural network to the maximum for a certain class.
 66
 67    Args:
 68        model (nn.Module): Torch model
 69        model_weights_path (str): Path to model saved weights.
 70        n_epochs (int, optional): Number of epochs to perform training. Defaults to 10.
 71        n_steps (int, optional): Number of steps within an epoch. Defaults to 100.
 72    """
 73
 74    # Create an image to start inverse training
 75    input = torch.rand((1, 28, 28))  # assume 28x28 image
 76    input.requires_grad = True  # set to make sure input can be adjusted
 77
 78    # Create embedding visualization
 79    model.load_state_dict(torch.load(model_weights_path))
 80
 81    # Freeze all parameters in the network
 82    for param in model.parameters():
 83        param.requires_grad = False
 84
 85    # Setup optimizer
 86    optimizer = SGD([input], lr=0.001)
 87    loss_fn = MSELoss()
 88
 89    target = torch.Tensor([10, -10, -10, -10, -10, -10, -10, -10, -10, -10])
 90    images = []
 91
 92    # Perform training epochs
 93    for ii in range(0, n_epochs):
 94        losses = []
 95        for jj in range(0, n_steps):
 96            pred, _ = model(input)
 97            loss = loss_fn(pred, target)
 98            optimizer.zero_grad()
 99            loss.backward()
100            optimizer.step()
101            losses.append(loss.item())
102
103        image = input.squeeze().clone().detach().numpy()
104        images.append(image)
105        print(f"Generative loss: {np.mean(losses)}")
106
107    # Draw images
108    for ii in range(1, len(images)):
109        ax = plt.subplot(1, len(images), ii)
110        ax.imshow(images[ii])