scripts/data_scaling_classifier.py

#!/usr/bin/env python3
"""DiffThinker C6: Data scaling proxy — ResNet-18 classifier on maze path-existence.
Binary classification: does a path exist from green start to red goal?
Shows accuracy improves with more training data, verifying the paper's scaling claim.
"""

import modal, json, os, sys, time, torch, torch.nn as nn
import numpy as np
from torch.utils.data import Dataset, DataLoader
from collections import deque

app = modal.App("diffthinker-data-scaling-classifier")

image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install("torch>=2.1.0", "torchvision", "Pillow", "numpy")
)


def bfs_path_exists(grid, start, goal):
    gs = grid.shape[0]
    q = deque([start])
    visited = {start}
    while q:
        r, c = q.popleft()
        if (r, c) == goal: return True
        for dr, dc in [(0,1),(0,-1),(1,0),(-1,0)]:
            nr, nc = r + dr, c + dc
            if 0 <= nr < gs and 0 <= nc < gs and grid[nr, nc] == 0 and (nr, nc) not in visited:
                visited.add((nr, nc))
                q.append((nr, nc))
    return False


def make_classification_data(num_samples, grid_size=8):
    """Generate maze images with binary labels: 1=path exists, 0=no path.
    Uses BFS to check solvability. Ensures ~50/50 split.
    """
    from PIL import Image, ImageDraw
    data = []
    targets = []
    solvable_count = 0
    target_solvable = num_samples // 2
    while len(data) < num_samples:
        grid = np.zeros((grid_size, grid_size), dtype=np.uint8)
        start = (0, np.random.randint(0, grid_size))
        goal = (grid_size - 1, np.random.randint(0, grid_size))
        wall_density = np.random.uniform(0.1, 0.35)
        for _ in range(int(grid_size * grid_size * wall_density)):
            wx, wy = np.random.randint(0, grid_size, 2)
            if (wx, wy) != start and (wx, wy) != goal:
                grid[wx, wy] = 1
        has_path = bfs_path_exists(grid, start, goal)
        if has_path and solvable_count >= target_solvable:
            continue  # already enough solvable
        if not has_path and len(data) - solvable_count >= num_samples - target_solvable:
            continue  # already enough unsolvable
        if has_path: solvable_count += 1
        img = Image.new("RGB", (64, 64), (255, 255, 255))
        draw = ImageDraw.Draw(img)
        cw = 64 // grid_size
        for r in range(grid_size):
            for c in range(grid_size):
                if grid[r, c] == 1:
                    draw.rectangle([c*cw, r*cw, (c+1)*cw, (r+1)*cw], fill=(100, 100, 100))
        draw.rectangle([start[1]*cw, start[0]*cw, (start[1]+1)*cw, (start[0]+1)*cw], fill=(0, 255, 0))
        draw.rectangle([goal[1]*cw, goal[0]*cw, (goal[1]+1)*cw, (goal[0]+1)*cw], fill=(255, 0, 0))
        data.append(img)
        targets.append(1 if has_path else 0)
    return data, targets


class ClassifierDS(Dataset):
    def __init__(self, images, labels):
        self.imgs = images
        self.lbls = labels
    def __len__(self):
        return len(self.imgs)
    def __getitem__(self, idx):
        img = torch.tensor(np.array(self.imgs[idx]).transpose(2,0,1), dtype=torch.float32) / 255.0
        return img, torch.tensor(self.lbls[idx], dtype=torch.long)


@app.function(image=image, gpu="T4", timeout=900)
def train_classifier(num_train, num_test=200, epochs=20):
    device = torch.device("cuda")
    gpu_name = torch.cuda.get_device_name(0)
    print(f"=== N={num_train} training samples | classifier | {epochs} epochs ===")

    from torchvision.models import resnet18, ResNet18_Weights
    train_imgs, train_lbls = make_classification_data(num_train)
    test_imgs, test_lbls = make_classification_data(num_test)
    print(f"  Train: {len(train_imgs)} ({sum(train_lbls)} solvable / {len(train_lbls)-sum(train_lbls)} blocked)")
    print(f"  Test:  {len(test_imgs)} ({sum(test_lbls)} solvable / {len(test_lbls)-sum(test_lbls)} blocked)")

    train_loader = DataLoader(ClassifierDS(train_imgs, train_lbls), batch_size=16, shuffle=True)
    test_ds = ClassifierDS(test_imgs, test_lbls)

    model = resnet18(weights=ResNet18_Weights.DEFAULT)
    model.fc = nn.Linear(model.fc.in_features, 2)
    model = model.to(device)
    total_p = sum(p.numel() for p in model.parameters())
    print(f"  Model params: {total_p:,}")

    optim = torch.optim.AdamW(model.parameters(), lr=1e-4)
    criterion = nn.CrossEntropyLoss()

    t_start = time.time()
    for ep in range(epochs):
        model.train()
        loss_sum = 0.0
        for imgs, lbls in train_loader:
            imgs, lbls = imgs.to(device), lbls.to(device)
            optim.zero_grad()
            out = model(imgs)
            loss = criterion(out, lbls)
            loss.backward()
            optim.step()
            loss_sum += loss.item()
        if (ep+1) % 5 == 0 or ep == epochs-1:
            print(f"  Ep {ep+1}/{epochs} loss={loss_sum/len(train_loader):.4f}")
    elapsed = time.time() - t_start

    model.eval()
    correct = 0
    with torch.no_grad():
        for img, lbl in test_ds:
            out = model(img.unsqueeze(0).to(device))
            pred = out.argmax(1).item()
            if pred == lbl: correct += 1

    acc = correct / num_test * 100
    print(f"  Acc: {acc:.1f}% ({correct}/{num_test}) | Time: {elapsed:.1f}s")
    return {"num_train": num_train, "params": total_p, "accuracy": round(acc,1),
            "epochs": epochs, "time_sec": round(elapsed,1), "gpu": gpu_name,
            "train_solvable": sum(train_lbls), "train_blocked": len(train_lbls)-sum(train_lbls)}


@app.function(image=image, gpu="T4", timeout=1800)
def classifier_sweep():
    ns = [20, 50, 100, 200, 400, 800, 1600]
    results = []
    for n in ns:
        r = train_classifier.remote(num_train=n)
        results.append(r)
        print(json.dumps(r))

    print("\n=== CLASSIFIER DATA SCALING RESULTS ===")
    print(f"{'N':>6} | {'Acc%':>6} | {'Time':>8}")
    print("-"*25)
    for r in results:
        print(f"{r['num_train']:>6} | {r['accuracy']:>5.1f}% | {r['time_sec']:>7.1f}s")
    return results


@app.local_entrypoint()
def main():
    print("="*60)
    print("C6: Data Scaling — ResNet-18 Classifier on T4")
    print("Paper: 20B MMDiT accuracy scales with training data")
    print("="*60)
    print("Task: Binary classification (path exists? on 8x8 maze images)")
    print("Model: ResNet-18 (~11M params) fine-tuned")
    print("Metric: Accuracy on 200 held-out test samples\n")
    r = classifier_sweep.remote()
    print(json.dumps(r, indent=2))