import torch

print(torch.version)

device = "cuda" if torch.cuda.is_available() else "cpu" print(device)

import torchvision from torchvision import datasets from torch import nn from torchvision.transforms import ToTensor from torchmetrics import Accuracy import matplotlib.pyplot as plt from torch.utils.data import DataLoader

train_data = datasets.MNIST( root="data", train=True, download=True, transform=ToTensor(), target_transform=None )

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

BATCH_SIZE = 32 train_dataloader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True )

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

)

class MNISTModelV0(nn.Module): def init(self, input_shape, hidden_units, output_shape): super().init() self.block_1 = nn.Sequential( nn.Conv2d(in_channels=input_shape, out_channels=hidden_units, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2) ) self.block_2 = nn.Sequential( nn.Conv2d(hidden_units, hidden_units, 3, padding=1), nn.ReLU(), nn.Conv2d(hidden_units, hidden_units, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2) ) self.classifier = nn.Sequential( nn.Flatten(), nn.Linear(in_features=hidden_units77, out_features=output_shape) )

def forward(self, x: torch.Tensor):
    x = self.block_1(x)
    x = self.block_2(x)
    x = self.classifier(x)
    return x

torch.manual_seed(42) mnist_model = MNISTModelV0(input_shape=1, hidden_units=10, output_shape=len(classes)).to(device) mnist_model

from tqdm.auto import tqdm loss_fn = nn.CrossEntropyLoss() accuracy_fn = Accuracy(task='multiclass', num_classes=len(classes)) optimizer= torch.optim.SGD(mnist_model.parameters(), lr=0.1)

def train_step(model, data_loader, loss_fn, optimizer, accuracy_fn, device): train_loss, train_acc = 0, 0 model.to(device) for batch, (X,y) in enumerate(data_loader): X, y = X.to(device), y.to(device) y_pred = model(X) loss = loss_fn(y_pred, y) train_loss += loss acc = accuracy_fn(y_pred.argmax(dim=1), y) train_acc += acc optimizer.zero_grad() loss.backward() optimizer.step() train_loss /= len(data_loader) train_acc /= len(data_loader) print(f"Train loss: {train_loss:.5f} | Train accuracy: {train_acc:.2f}%")

def test_step(model, data_loader, loss_fn, accuracy_fn, device): test_loss, test_acc = 0, 0 model.to(device) model.eval() with torch.inference_mode(): for (X,y) in data_loader: X, y = X.to(device), y.to(device) y_pred = model(X) test_loss += loss_fn(y_pred, y) test_acc += accuracy_fn(y_pred.argmax(dim=1), y) test_loss /= len(data_loader) test_acc /= len(data_loader) print(f"Test loss: {test_loss:.5f} | Test accuracy: {test_acc:.2f}%")

epochs = 5

for epoch in tqdm(range(epochs)): print(f"Epoch: {epoch}\n---------") train_step(data_loader=train_dataloader, model=mnist_model, loss_fn=loss_fn, optimizer=optimizer, accuracy_fn=accuracy_fn, device=device ) test_step(data_loader=test_dataloader, model=mnist_model, loss_fn=loss_fn, accuracy_fn=accuracy_fn, device=device )