model.py

# Import PyTorch
import torch
from torch import nn

# Import torchvision
import torchvision
from torchvision import datasets
from torchvision import transforms
from torchvision.transforms import ToTensor
from torch.utils.data import DataLoader

# Import matplotlib
import matplotlib.pyplot as plt


# Check versions
print(torch.__version__)
print(torchvision.__version__)

# Setup training data
train_data = datasets.MNIST(
    root="./model/data",
    train=True,
    download=True,
    transform=ToTensor(),
    target_transform=None
)

test_data = datasets.MNIST(
    root="./model/data",
    train=False,
    download=True,
    transform=ToTensor(),
    target_transform=None
)

class_names = train_data.classes

# print(len(train_data))
# print(len(test_data))

# image, label = train_data[0]
# print(f"Image shape: {image.shape}")
# plt.imshow(image.squeeze())
# plt.title(label)
# plt.axis(False)
# plt.show()

# Plot more images
# torch.manual_seed(42)
# fig = plt.figure(figsize=(9, 9))
# rows, cols = 4, 4
# for i in range(1, rows*cols+1):
#   random_idx = torch.randint(0, len(train_data), size=[1]).item()
#   img, label = train_data[random_idx]
#   fig.add_subplot(rows, cols, i)
#   plt.imshow(img.squeeze(), cmap="gray")
#   plt.title(class_names[label])
#   plt.axis(False)

# Setup the batch size
BATCH_SIZE = 32

# turn datasets into iterables
train_dataloader = DataLoader(dataset=train_data,
                              batch_size=BATCH_SIZE,
                              shuffle=True)

test_dataloader = DataLoader(dataset=test_data,
                             batch_size=BATCH_SIZE,
                             shuffle=False)

class FashionMNISTModelV0(nn.Module):
  def __init__(self,
               input_shape: int,
               hidden_units: int,
               output_shape: int):
    super().__init__()
    self.layer_stack = nn.Sequential(
        nn.Flatten(),
        nn.Linear(in_features=input_shape,
                  out_features=hidden_units),
        nn.Linear(in_features=hidden_units,
                  out_features=output_shape)
    )

  def forward(self, x):
    return self.layer_stack(x)
  

# Setup model with input parameters
model_0 = FashionMNISTModelV0(
    input_shape=784, # this is 28 x 28
    hidden_units=10,
    output_shape=len(class_names) # one for every class
)
model_0.to("cpu")

# Setup loss fn and optimizer
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(
    params=model_0.parameters(),
    lr=0.1
)

# Set the seet and start the timer
torch.manual_seed(42)

# Set the number of epochs
epochs = 3

# Create training and test loop
for epoch in range(epochs):
  print(f"Epoch: {epoch}\n------")
  ### Training
  train_loss = 0
  # Add a loop to loop through the training batches
  for batch, (X, y) in enumerate(train_dataloader):
    model_0.train()
    # 1. Forward pass
    y_pred = model_0(X)

    # 2. Calculate the loss
    loss = loss_fn(y_pred, y)
    train_loss += loss

    # 3. Optimizer zero grad
    optimizer.zero_grad()

    # 4. Loss backward
    loss.backward()

    # 5. Optimizer step
    optimizer.step()

    # Print out what's happening
    if batch % 400 == 0:
      print(f"Looked at {batch * len(X)}/{len(train_dataloader.dataset)}")

  # Divide total train loss by length of train dataloader
  train_loss /= len(train_dataloader)

  ### Testing
  test_loss, test_acc = 0, 0
  model_0.eval()
  with torch.inference_mode():
    for X_test, y_test in test_dataloader:
        # 1. Forward pass
        test_pred = model_0(X_test)  

        # 2. Calculate the loss
        test_loss += loss_fn(test_pred, y_test)

    # Calculate the test loss and accuracy average per batch
    test_loss /= len(test_dataloader)
    test_acc /= len(test_dataloader)

  # Print out whats happening
  print(f"\nTrain loss: {train_loss:.4f} | Test loss : {test_loss:.4f} | Test acc: {test_acc:.4f}")

# Saving our PyTorch model
from pathlib import Path

# 1. Create models directory
MODEL_PATH = Path("model")
MODEL_PATH.mkdir(parents=True, exist_ok=True)

# 2. Create model save path
MODEL_NAME = "model.pth"
MODEL_SAVE_PATH = MODEL_PATH / MODEL_NAME

# 3. Save the model state dict
print(f"Saving model to: {MODEL_SAVE_PATH}")
torch.save(obj=model_0.state_dict(),
           f=MODEL_SAVE_PATH)