PyTorch training¶

"""Example MNIST training script"""

import pathlib
import tempfile

import progressbar
import torch
from torch import nn
from torchvision import datasets, transforms
import typer

app = typer.Typer()


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = nn.functional.relu(x)
        x = self.conv2(x)
        x = nn.functional.relu(x)
        x = nn.functional.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = nn.functional.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        return nn.functional.log_softmax(x, dim=1)


@app.command()
def prepare_dataset(output: pathlib.Path):
    with tempfile.TemporaryDirectory() as download_path:
        print("Downloading MNIST dataset...")
        transform = transforms.Compose(
            [
                transforms.ToTensor(),
                transforms.Normalize((0.1307,), (0.3081,)),
            ]
        )
        train_tune = datasets.MNIST(download_path, download=True, transform=transform)
        test = datasets.MNIST(download_path, train=False, transform=transform)

        print("Preprocessing and saving train/tune/test splits...")
        train, tune = torch.utils.data.random_split(train_tune, [0.8, 0.2])
        torch.save(list(train), output / "train.pt")
        torch.save(list(tune), output / "tune.pt")
        torch.save(list(test), output / "test.pt")


@app.command()
def train(
    dataset: pathlib.Path,
    model_output: pathlib.Path,
    batch_size: int = 64,
    epochs: int = 20,
    learning_rate: float = 1.0,
):
    if torch.accelerator.is_available():
        device = torch.accelerator.current_accelerator()
    else:
        device = torch.device("cpu")

    train_loader = torch.utils.data.DataLoader(
        torch.load(dataset / "train.pt"),
        batch_size=batch_size,
        shuffle=True,
    )
    tune_loader = torch.utils.data.DataLoader(
        torch.load(dataset / "tune.pt"),
        batch_size=batch_size,
    )

    model = Model().to(device)
    optimizer = torch.optim.Adadelta(model.parameters(), lr=learning_rate)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.7)

    for epoch in range(1, epochs + 1):
        model.train()
        bar = progressbar.progressbar(
            train_loader,
            desc=f"Epoch {epoch}",
            fast=False,
            min_poll_interval=2,
        )
        for data, target in bar:
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            output = model(data)
            loss = nn.functional.nll_loss(output, target)
            loss.backward()
            optimizer.step()
        scheduler.step()

        model.eval()
        loss = 0
        correct = 0
        with torch.no_grad():
            for data, target in tune_loader:
                data, target = data.to(device), target.to(device)
                output = model(data)
                loss += nn.functional.nll_loss(output, target, reduction="sum").item()
                pred = output.argmax(dim=1, keepdim=True)
                correct += pred.eq(target.view_as(pred)).sum().item()

            loss /= len(tune_loader.dataset)
            accuracy = 100 * correct / len(tune_loader.dataset)
            print(f"Tune set: average loss: {loss:.4f}, accuracy: {accuracy:.1f}%")

    torch.save(model.cpu().state_dict(), model_output)


@app.command()
def evaluate(model_path: pathlib.Path, dataset: pathlib.Path):
    test_loader = torch.utils.data.DataLoader(torch.load(dataset / "test.pt"))

    model = Model()
    model.load_state_dict(torch.load(model_path, weights_only=True))
    model.eval()

    correct = 0
    for data, target in test_loader:
        output = model(data)
        pred = output.argmax(dim=1, keepdim=True)
        correct += pred.eq(target.view_as(pred)).sum().item()

    accuracy = 100 * correct / len(test_loader.dataset)
    print(f"Test set accuracy: {accuracy:.1f}%")


if __name__ == "__main__":
    app()