Deep Learning and AI Superhero

Practical Exercises Chapter 4

Exercise 1: Saving and Loading a Model’s state_dict

Task: Define a simple neural network, train it on the MNIST dataset, save the model’s state_dict, and then load the saved state_dict to continue training from where it left off.

Solution:

import torch

import torch.nn as nn

import torch.optim as optim

from torchvision import datasets, transforms

from torch.utils.data import DataLoader

# Define a simple neural network

class SimpleNN(nn.Module):

def __init__(self):

super(SimpleNN, self).__init__()

self.fc1 = nn.Linear(784, 128)

self.fc2 = nn.Linear(128, 64)

self.fc3 = nn.Linear(64, 10)

def forward(self, x):

x = x.view(-1, 784)

x = torch.relu(self.fc1(x))

x = torch.relu(self.fc2(x))

return self.fc3(x)

# Instantiate the ...