Spaces:
Running on Zero
Running on Zero
File size: 4,059 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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | """Render Sudoku grids as images of MNIST digits and build training/eval batches.
Image: 1 x 144 x 144 in [-1, 1] (9 x 9 cells of 16 px, the paper's training resolution).
Conditioning: the solution image with non-clue cells blacked out, plus a pixel mask of clue cells.
"""
import numpy as np
import torch
import torch.nn.functional as F
from . import sudoku
CELL = 16
IMG = 9 * CELL
class DigitBank:
"""MNIST digits resized to CELL x CELL, grouped by class: images[d] is (N_d, CELL, CELL) in [-1, 1]."""
def __init__(self, train=True, data_dir="data", images=None, labels=None):
if images is None:
from torchvision.datasets import MNIST
ds = MNIST(data_dir, train=train, download=True)
images, labels = ds.data, ds.targets
x = F.interpolate(images[:, None].float() / 255.0, size=(CELL, CELL), mode="area")[:, 0]
x = x * 2 - 1
n = min(int((labels == d).sum()) for d in range(10))
self.images = torch.stack([x[labels == d][:n] for d in range(10)]) # (10, n, CELL, CELL)
def to(self, device):
self.images = self.images.to(device)
return self
def render(grids, bank, generator=None):
"""grids: (B, 9, 9) long tensor of digits 1-9 (0 renders as black). Returns (B, 1, IMG, IMG)."""
B = grids.shape[0]
flat = grids.reshape(B, 81).to(bank.images.device)
idx = torch.randint(bank.images.shape[1], (B, 81), generator=generator, device="cpu").to(flat.device)
cells = bank.images[flat, idx] # (B, 81, CELL, CELL)
cells = torch.where((flat == 0)[..., None, None], torch.full_like(cells, -1.0), cells)
return cells.reshape(B, 9, 9, CELL, CELL).permute(0, 1, 3, 2, 4).reshape(B, 1, IMG, IMG)
def cells(images):
"""Inverse of the layout in render: (B, 1, IMG, IMG) -> (B, 81, CELL, CELL)."""
B = images.shape[0]
return images.reshape(B, 9, CELL, 9, CELL).permute(0, 1, 3, 2, 4).reshape(B, 81, CELL, CELL)
def pixel_mask(puzzles):
"""(B, 9, 9) puzzles -> (B, 1, IMG, IMG) float mask, 1 on clue cells."""
m = (puzzles > 0).float()
return m.repeat_interleave(CELL, 1).repeat_interleave(CELL, 2)[:, None]
def make_batch(puzzles, solutions, bank, generator=None):
"""Render a batch. Clue pixels in cond are the same digit images as in x0."""
puzzles, solutions = torch.as_tensor(puzzles).long(), torch.as_tensor(solutions).long()
x0 = render(solutions, bank, generator)
mask = pixel_mask(puzzles).to(x0.device)
cond = torch.where(mask > 0, x0, torch.full_like(x0, -1.0))
return dict(x0=x0, cond=cond, mask=mask, puzzles=puzzles, solutions=solutions)
class PuzzleSource:
"""Samples (puzzles, solutions) for 'extreme' (TRM csv), 'hard' (0-26 random clues of a real solution)
or 'solutions' (no clues; Painter data). Solutions always come from the sudoku-extreme split, so train and
test grids are disjoint; training samples get a random validity-preserving augmentation."""
def __init__(self, mode, split="train", data_dir="data", seed=0):
if mode not in ("extreme", "hard", "solutions"):
raise ValueError(mode)
self.mode, self.split, self.rng = mode, split, np.random.default_rng(seed)
self.puzzles, self.solutions = sudoku.load_extreme(split, data_dir)
def sample(self, n):
i = self.rng.integers(len(self.solutions), size=n)
p, s = sudoku.augment(self.rng, self.puzzles[i], self.solutions[i])
if self.mode == "hard":
p = np.stack([sudoku.hard_puzzle(self.rng, g)[0] for g in s])
elif self.mode == "solutions":
p = np.zeros_like(p)
return p, s
def fixed(self, n, seed=1234):
"""A deterministic evaluation set (no augmentation)."""
rng = np.random.default_rng(seed)
i = rng.choice(len(self.solutions), size=n, replace=False)
p, s = self.puzzles[i].astype(np.int64), self.solutions[i].astype(np.int64)
if self.mode == "hard":
p = np.stack([sudoku.hard_puzzle(rng, g)[0] for g in s])
return p, s
|