Spaces:
Running on Zero
Running on Zero
File size: 2,843 Bytes
65635a9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | """Small MNIST CNN on 16x16 cells, used to read generated Sudoku images back into grids."""
import torch
import torch.nn as nn
import torch.nn.functional as F
from .data import cells
def conv_bn(cin, cout):
return [nn.Conv2d(cin, cout, 3, padding=1), nn.BatchNorm2d(cout), nn.ReLU()]
class DigitCNN(nn.Module):
def __init__(self):
super().__init__()
self.f = nn.Sequential(
*conv_bn(1, 64), *conv_bn(64, 64), nn.MaxPool2d(2),
*conv_bn(64, 128), *conv_bn(128, 128), nn.MaxPool2d(2),
*conv_bn(128, 256), nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Dropout(0.3), nn.Linear(256, 10))
def forward(self, x):
return self.f(x)
@torch.no_grad()
def read_grids(clf, images):
"""(B, 1, 144, 144) images -> (B, 9, 9) predicted digits in 1..9."""
c = cells(images)
B = c.shape[0]
logits = clf(c.reshape(B * 81, 1, *c.shape[-2:]))
return (logits[:, 1:].argmax(-1) + 1).reshape(B, 9, 9)
def augment(x):
"""Random small rotation, scale and shift (cell digits are not perfectly centred)."""
B = len(x)
a = (torch.rand(B, device=x.device) - 0.5) * 0.4
s = 1 + (torch.rand(B, device=x.device) - 0.5) * 0.2
theta = torch.zeros(B, 2, 3, device=x.device)
theta[:, 0, 0], theta[:, 0, 1] = s * a.cos(), -s * a.sin()
theta[:, 1, 0], theta[:, 1, 1] = s * a.sin(), s * a.cos()
theta[:, :, 2] = (torch.rand(B, 2, device=x.device) - 0.5) * 0.25
grid = F.affine_grid(theta, x.shape, align_corners=False)
return F.grid_sample(x + 1, grid, align_corners=False) - 1 # pad with background (-1)
def train(train_bank, test_bank, device, epochs=10, batch=128, lr=2e-3, seed=0):
"""Train on DigitBank tensors (all 10 classes). Returns (model, test accuracy)."""
torch.manual_seed(seed)
def flat(bank):
x = bank.images
y = torch.arange(10, device=x.device)[:, None].expand(x.shape[:2])
return x.reshape(-1, 1, *x.shape[-2:]).to(device), y.reshape(-1).to(device)
x, y = flat(train_bank)
xt, yt = flat(test_bank)
model = DigitCNN().to(device)
opt = torch.optim.AdamW(model.parameters(), lr, weight_decay=5e-4)
sched = torch.optim.lr_scheduler.OneCycleLR(opt, lr, total_steps=epochs * ((len(x) + batch - 1) // batch))
for ep in range(epochs):
model.train()
perm = torch.randperm(len(x), device=device)
for i in range(0, len(x), batch):
idx = perm[i:i + batch]
loss = F.cross_entropy(model(augment(x[idx])), y[idx], label_smoothing=0.05)
opt.zero_grad()
loss.backward()
opt.step()
sched.step()
model.eval()
with torch.no_grad():
acc = sum((model(xt[i:i + 2048]).argmax(-1) == yt[i:i + 2048]).sum().item() for i in range(0, len(xt), 2048)) / len(xt)
return model, acc
|