Spaces:
Running on Zero
Running on Zero
| """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) | |
| 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 | |