PyTorch Lightning: MNIST

pip install torch torchvision pytorch-lightning

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import pytorch_lightning as pl

SEED = 42
pl.seed_everything(SEED, verbose=False)
print(f"KVERITAS_INPUT src=seed:{SEED}", flush=True)

print("KVERITAS_PHASE name=data", flush=True)
tfm = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
train_ds = datasets.MNIST("./data", train=True, download=True, transform=tfm)
test_ds = datasets.MNIST("./data", train=False, download=True, transform=tfm)
train_loader = DataLoader(train_ds, batch_size=128, shuffle=True)
test_loader = DataLoader(test_ds, batch_size=512)


class LitMLP(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Flatten(), nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10))

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

    def training_step(self, batch, _):
        x, y = batch
        loss = F.cross_entropy(self(x), y)
        self.log("train_loss", loss, on_epoch=True, prog_bar=False)
        return loss

    def test_step(self, batch, _):
        x, y = batch
        acc = (self(x).argmax(1) == y).float().mean()
        self.log("test_accuracy", acc)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)


model = LitMLP()
n_params = sum(p.numel() for p in model.parameters())
print(f"KVERITAS_MODEL params={n_params} arch=mlp precision=fp32", flush=True)
print(f"KVERITAS_WORKLOAD dataset_size={len(train_ds)} epochs=2 batch_size=128", flush=True)


class KVeritas(pl.Callback):
    def on_train_epoch_end(self, trainer, module):
        loss = trainer.callback_metrics.get("train_loss_epoch") or trainer.callback_metrics.get("loss")
        if loss is not None:
            print(f"KVERITAS_METRIC name=train_loss value={float(loss):.4f} step={trainer.current_epoch}", flush=True)


print("KVERITAS_PHASE name=train", flush=True)
trainer = pl.Trainer(max_epochs=2, accelerator="auto", devices=1, logger=False,
                     enable_checkpointing=False, enable_progress_bar=False, callbacks=[KVeritas()])
trainer.fit(model, train_loader)

print("KVERITAS_PHASE name=evaluate", flush=True)
result = trainer.test(model, test_loader, verbose=False)
acc = float(result[0]["test_accuracy"])
print(f"KVERITAS_METRIC name=test_accuracy value={acc:.4f}", flush=True)
print(f"KVERITAS_CLAIM metric=test_accuracy value={acc:.4f}", flush=True)
print(f"final test_accuracy={acc:.4f}", flush=True)

kveritas init
kveritas run -- python train.py
kveritas seal --output report.pdf
kveritas verify report.pdf