PyTorch: KMNIST
A small convolutional network on KMNIST (Kuzushiji-MNIST), a drop-in but harder replacement for MNIST made of cursive Japanese characters, downloaded through torchvision.
Install
The framework, plus the K-Veritas CLI (see the Overview for install).
pip install torch torchvision
train.py
The complete script. The KVERITAS_ lines are the only additions to an ordinary training script; everything else is standard PyTorch.
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
SEED = 42
torch.manual_seed(SEED)
print(f"KVERITAS_INPUT src=seed:{SEED}", flush=True)
device = "cuda" if torch.cuda.is_available() else "cpu"
print("KVERITAS_PHASE name=data", flush=True)
tfm = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1918,), (0.3483,))])
train_ds = datasets.KMNIST("./data", train=True, download=True, transform=tfm)
test_ds = datasets.KMNIST("./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 Net(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 32, 3, padding=1)
self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
self.drop = nn.Dropout(0.25)
self.fc1 = nn.Linear(64 * 7 * 7, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = F.max_pool2d(F.relu(self.conv1(x)), 2)
x = F.max_pool2d(F.relu(self.conv2(x)), 2)
x = self.drop(x.flatten(1))
x = F.relu(self.fc1(x))
return self.fc2(x)
model = Net().to(device)
n_params = sum(p.numel() for p in model.parameters())
print(f"KVERITAS_MODEL params={n_params} arch=cnn precision=fp32", flush=True)
print(f"KVERITAS_WORKLOAD dataset_size={len(train_ds)} epochs=4 batch_size=128", flush=True)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
def accuracy(loader):
model.eval()
correct = total = 0
with torch.no_grad():
for x, y in loader:
x, y = x.to(device), y.to(device)
correct += (model(x).argmax(1) == y).sum().item()
total += y.size(0)
return correct / total
print("KVERITAS_PHASE name=train", flush=True)
for epoch in range(1, 5):
model.train()
running = 0.0
for x, y in train_loader:
x, y = x.to(device), y.to(device)
opt.zero_grad()
loss = F.cross_entropy(model(x), y)
loss.backward()
opt.step()
running += loss.item() * y.size(0)
train_loss = running / len(train_ds)
test_acc = accuracy(test_loader)
print(f"KVERITAS_METRIC name=train_loss value={train_loss:.4f} step={epoch}", flush=True)
print(f"KVERITAS_METRIC name=test_accuracy value={test_acc:.4f} step={epoch}", flush=True)
print(f"epoch {epoch}/4 train_loss={train_loss:.4f} test_acc={test_acc:.4f}", flush=True)
print("KVERITAS_PHASE name=evaluate", flush=True)
final_acc = accuracy(test_loader)
print(f"KVERITAS_CLAIM metric=test_accuracy value={final_acc:.4f}", flush=True)
print(f"final test_accuracy={final_acc:.4f}", flush=True)Run and seal
Wrap the script, seal a signed report, and verify it offline.
kveritas init kveritas run -- python train.py kveritas seal --output report.pdf kveritas verify report.pdf
What it produced here
4 epochs reached about 95.4% test accuracy on Kuzushiji-MNIST, a harder handwriting dataset than MNIST.