Section 8 - DL HW 1

import json

class Network(nn.Module):
    def __init__(self, image_width, image_height):
        super().__init__()
        self.image_width = image_width
        self.image_height = image_height

        self.conv1 = nn.Conv2d(3, 32, 3)
        self.conv2 = nn.Conv2d(32, 16, 3)
        self.pool = nn.MaxPool2d(2, 2)

        self.fc1 = nn.Linear(16 * 62 * 62, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 62 * 62)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# Create the model
model_custom = Network(256, 256)

# If a GPU is available, make the model use it
model_custom = model_custom.to(device)

criterion = nn.CrossEntropyLoss()
optimizer_ft = optim.SGD(model_custom.parameters(), lr=0.03, momentum=0.4)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)
num_epochs = 20

# Train the model
model_custom, performance_custom = train_model(model_custom, 
                                               dataloaders,
                                               criterion,
                                               optimizer_ft,
                                               exp_lr_scheduler,
                                               dataset_sizes,
                                               num_epochs=num_epochs)

# Save the model and performance for later
model_file = os.path.join(models_path, 'model_custom.pth')
torch.save(model_custom.state_dict(), model_file)

with open(os.path.join(metrics_path, 'metrics_custom.json'), 'w') as f:
    json.dump(performance_custom, f)
    
#######################
# Plot the convergence
#######################

# Create subplots for Accuracy and Loss plots
fig, axs = plt.subplots(1, 2, figsize=(10, 6))
metrics = ['acc', 'loss']
titles = ['Accuracy as a function of epochs', 'Cross Entropy Loss as a function of epochs']

# Plot the model's convergence
for i, metric in enumerate(metrics):
  graph = axs[i]
  graph.set_title(titles[i])
  graph.set_xlabel('Epochs')

  for s in ['train', 'val']:
    graph.plot(np.arange(num_epochs), performance_custom[s][metric], label=s)
  graph.legend(['train', 'val'])