MNIST Multi-Layer Perceptron with Validation

INPUT_NODES = 28 * 28
HIDDEN_NODES_1 = HIDDEN_NODES_2 = 512
DROPOUT = 0.2
CLASSES = 10

class MultiLayerPerceptron(nn.Module):
    """A Multi-Layer Perceptron

    This is a network with 2 hidden layers
    """
    def __init__(self):
        super().__init__()        
        self.fully_connected_layer_1 = nn.Linear(INPUT_NODES, HIDDEN_NODES_1)
        self.fully_connected_layer_2 = nn.Linear(HIDDEN_NODES_1, HIDDEN_NODES_2)
        self.output = nn.Linear(HIDDEN_NODES_2, CLASSES)
        self.dropout = nn.Dropout(p=DROPOUT)
        return

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """One feed-forward through the network

       Args:
        x: a 28 x 28 tensor

       Returns:
        tensor: output of the network without activation
       """
        # flatten image input
        x = x.view(-1, INPUT_NODES)

        x = self.dropout(F.relu(self.fully_connected_layer_1(x)))
        x = self.dropout(F.relu(self.fully_connected_layer_2(x)))        
        return self.output(x)

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=LEARNING_RATE)

# number of epochs to train the model
EPOCHS = 50
SAVED_MODEL= 'multilayer_perceptron.pt'

 def process_batch(model: nn.Module, data: torch.Tensor, target: torch.Tensor,
                   device: str) -> Outcome:
     """process one batch of the data

     Args:
      model: model to predict target
      data: data to use to predict target
      target: what we're trying to predict
      device: cpu or gpu

     Returns:
      outcome: loss and correct count
     """
     data, target = data.to(device), target.to(device)
     output = model(data)
     loss = criterion(output, target)
     _, predicted = torch.max(output.data, 1)
     return loss, (predicted == target).sum().item()

def train(model: nn.Module,
          batches: torch.utils.data.DataLoader,
          device: str,
) -> Outcome:
    """Perform one forward pass through the batches

    Args:
     model: thing to train
     batches: batch-iterator of training data
     device: cpu or cuda device

    Returns:
     outcome: cumulative loss, accuracy for the batches
    """
    total_loss = 0.0
    count = 0
    total_correct = 0
    model.train()
    for data, target in batches:
        optimizer.zero_grad()
        loss, correct = process_batch(model, data, target, device)
        count += target.size(0)
        total_correct += correct
        total_loss += loss
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * data.size(0)
    return total_loss, total_correct/count

def validate(model: nn.Module, batches: torch.utils.data.DataLoader,
             device: str) -> Outcome:
    """Calculate the loss for the model

    Args:
     model: the model to validate
     batches: the batch-iterator of validation data
     device: cuda or cpu

    Returns:
     Outcome: Cumulative loss, Accuracy over batches
    """
    model.eval()
    total_loss = 0.0
    total_correct = 0
    count = 0
    for data, target in batches:
        loss, correct = process_batch(model, data, target, device)
        count += target.size(0)
        total_correct += correct
        total_loss += loss.item() * data.size(0)
    return total_loss, total_correct/count

# initialize tracker for minimum validation loss
lowest_validation_loss = numpy.Inf
training_losses = []
validation_losses = []
training_accuracies = []
validation_accuracies = []
for epoch in range(1, EPOCHS + 1):
    loss, accuracy = train(model, training_batches, device)
    training_losses.append(loss)
    mean_training_loss = loss/len(training_batches.dataset)
    training_accuracies.append(accuracy)

    loss, accuracy = validate(model, validation_batches, device)
    validation_losses.append(loss)
    mean_validation_loss = loss/len(validation_batches.dataset)
    validation_accuracies.append(accuracy)

    if mean_validation_loss <= lowest_validation_loss:
        print('Epoch {}: Validation loss decreased ({:.6f} --> {:.6f}).  Saving model ...'.format(
            epoch,
            lowest_validation_loss,
            mean_validation_loss))
        torch.save(model.state_dict(), SAVED_MODEL)
        lowest_validation_loss = mean_validation_loss

x = list(range(len(training_losses)))
figure, axe = pyplot.subplots()
figure.suptitle("Loss Per Batch", weight="bold")
axe.plot(x, training_losses, label="Training")
axe.plot(x, validation_losses, label="Validation")
legend = axe.legend()

model.load_state_dict(torch.load(SAVED_MODEL))

test_loss = 0.0
class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))

model.eval()

for data, target in test_batches:
    output = model(data)
    data, target = data.to(device), target.to(device)
    # calculate the loss
    loss = criterion(output, target)
    # update test loss 
    test_loss += loss.item()*data.size(0)
    # convert output probabilities to predicted class
    _, pred = torch.max(output, 1)
    # compare predictions to true label
    correct = numpy.squeeze(pred.eq(target.data.view_as(pred)))
    # calculate test accuracy for each object class
    for i in range(BATCH_SIZE):
        label = target.data[i]
        class_correct[label] += correct[i].item()
        class_total[label] += 1

# calculate and print avg test loss
test_loss = test_loss/len(test_batches.dataset)
print('Test Loss: {:.6f}\n'.format(test_loss))

for i in range(10):
    if class_total[i] > 0:
        print('Test Accuracy of %5s: %2d%% (%2d/%2d)' % (
            str(i), 100 * class_correct[i] / class_total[i],
            numpy.sum(class_correct[i]), numpy.sum(class_total[i])))
    else:
        print('Test Accuracy of %5s: N/A (no training examples)' % (classes[i]))

images, labels = iter(test_batches).next()
# matplotlib doesn't like the CUDA and the model doesn't like the CPU... too bad for the model.
model.to("cpu")
output = model(images)

_, preds = torch.max(output, 1)
# prep images for display
images = images.numpy()

# plot the images in the batch, along with predicted and true labels
figure = pyplot.figure(figsize=(25, 4))
title = figure.suptitle("Test Predictions", weight="bold", position=(0.5, 1.3))

for index in numpy.arange(20):
    ax = figure.add_subplot(2, 20/2, index+1, xticks=[], yticks=[])
    ax.imshow(numpy.squeeze(images[index]), cmap='gray')
    ax.set_title("{} ({})".format(str(preds[index].item()), str(labels[index].item())),
                 color=("green" if preds[index]==labels[index] else "red"))
figure.tight_layout()

 class Trainer:
     """Train-test-validate the model

     Args:
      train: training batches
      validate: validation batches
      test: testing batches
      epochs: number of times to repeat training over the batches
      model_filename: name to save the hyperparameters of best model
      learning_rate: how much to update the weights
     """
     def __init__(self, train: torch.utils.data.DataLoader,
                  validate: torch.utils.data.DataLoader,
                  test: torch.utils.data.DataLoader,
                  epochs: int=50,
                  model_filename: str="multilayer_perceptron.pth",
                  learning_rate=0.01) -> None:
         self.training_batches = train
         self.validation_batches = validate
         self.test_batches = test
         self.epochs = epochs
         self.save_as = model_filename
         self.learning_rate = learning_rate
         self._model = None
         self._criterion = None
         self._optimizer = None
         self._device = None
         self.validation_losses = []
         self.training_losses = []
         self.validation_accuracies = []
         self.training_accuracies = []
         self.best_parameters = None
         return

     @property
     def model(self):
         """The Multi-Layer Perceptron"""
         if self._model is None:
             self._model = model = MultiLayerPerceptron()
             self._model.to(self.device)
         return self._model

     @property
     def criterion(self):
         """The Loss Measurer"""
         if self._criterion is None:
             self._criterion = nn.CrossEntropyLoss()
         return self._criterion

     @property
     def optimizer(self):
         """The gradient descent"""
         if self._optimizer is None:
             self._optimizer = torch.optim.SGD(self.model.parameters(),
                                               lr=self.learning_rate)
         return self._optimizer

     @property
     def device(self):
         """The CPU or GPU"""
         if self._device is None:
             self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         return self._device


     def process_batch(self, data: torch.Tensor, target: torch.Tensor) -> Outcome:
         """process one batch of the data

        Args:
         data: data to use to predict target
         target: what we're trying to predict
         device: cpu or gpu

        Returns:
         outcome: loss and correct count
        """
         data, target = data.to(self.device), target.to(self.device)
         output = self.model(data)
         loss = self.criterion(output, target)
         _, predicted = torch.max(output.data, 1)
         return loss, (predicted == target).sum().item()

     def train(self) -> Outcome:
         """Perform one forward pass through the batches

        Returns:
         outcome: cumulative loss, accuracy for the batches
        """
         total_loss = 0.0
         count = 0
         total_correct = 0
         self.model.train()
         for data, target in self.training_batches:
             self.optimizer.zero_grad()
             loss, correct = self.process_batch(data, target)
             count += target.size(0)
             total_correct += correct
             total_loss += loss
             loss.backward()
             self.optimizer.step()
             total_loss += loss.item() * data.size(0)
             del loss
         return float(total_loss), float(total_correct/count)

     def validate(self) -> Outcome:
         """Calculate the loss for the model

        Returns:
         Outcome: Cumulative loss, Accuracy over batches
        """
         self.model.eval()
         total_loss = 0.0
         total_correct = 0
         count = 0
         for data, target in self.validation_batches:
             loss, correct = self.process_batch(data, target)
             count += target.size(0)
             total_correct += correct
             total_loss += loss.item() * data.size(0)
             del loss
         return float(total_loss), float(total_correct/count)

     def run_training(self) -> None:
         """Runs the training and validation"""
         lowest_validation_loss = numpy.Inf
         for epoch in range(1, self.epochs + 1):
             gc.collect()
             loss, accuracy = self.train()
             self.training_losses.append(loss)
             mean_training_loss = loss/len(self.training_batches.dataset)
             self.training_accuracies.append(accuracy)
             loss, accuracy = self.validate()
             self.validation_losses.append(loss)
             mean_validation_loss = loss/len(self.validation_batches.dataset)
             self.validation_accuracies.append(accuracy)

             if mean_validation_loss <= lowest_validation_loss:
                 print('Epoch {}: Validation loss decreased ({:.6f} --> {:.6f}).  Saving model ...'.format(
                     epoch,
                     lowest_validation_loss,
                     mean_validation_loss))
                 self.best_parameters = self.model.state_dict()
                 torch.save(self.best_parameters, self.save_as)
                 lowest_validation_loss = mean_validation_loss
         return

     def test(self):
         """Test Our Model"""
         if self.best_parameters is None:
             raise Exception("call ``run_training`` or set ``best_parameters")
         self.model.load_state_dict(self.best_parameters)
         test_loss = 0.0
         digits = 10
         class_correct = [0.0] * digits
         class_total = [0.0] * digits
         self.model.eval()

         for data, target in self.test_batches:
             output = self.model(data)
             data, target = data.to(device), target.to(device)
             loss = self.criterion(output, target)
             test_loss += loss.item() * data.size(0)

             _, predictions = torch.max(output, 1)
             correct = numpy.squeeze(predictions.eq(
                 target.data.view_as(predictions)))
             # calculate test accuracy for each object class
             for i in range(data.size(0)):
                 label = target.data[i]
                 class_correct[label] += correct[i].item()
                 class_total[label] += 1

         # calculate and print avg test loss
         test_loss = test_loss/len(self.test_batches.dataset)
         print('Test Loss: {:.6f}\n'.format(test_loss))

         for digit in range(10):
             if class_total[digit] > 0:
                 print('Test Accuracy of %5s: %2d%% (%2d/%2d)' % (
                     str(digit), 100 * class_correct[digit] / class_total[digit],
                     numpy.sum(class_correct[digit]), numpy.sum(class_total[digit])))
             else:
                 print('Test Accuracy of %5s: N/A (no training examples)' % (classes[digit]))
         return

 trainer = Trainer(training_batches, validation_batches, test_batches)
 trainer.run_training()

x = list(range(len(trainer.training_accuracies)))
figure, axe = pyplot.subplots()
figure.suptitle("Model Accuracy", weight="bold")
axe.plot(x, trainer.training_accuracies, label="Training")
axe.plot(x, trainer.validation_accuracies, label="Validation")
legend = axe.legend()

x = list(range(len(trainer.training_losses)))
figure, axe = pyplot.subplots()
figure.suptitle("Loss Per Batch", weight="bold")
axe.plot(x, trainer.training_losses, label="Training")
axe.plot(x, trainer.validation_losses, label="Validation")
legend = axe.legend()