scripts/hf_jobs/hf_datasize.py

# /// script
# requires-python = ">=3.12"
# dependencies = ["torch>=2.5.0", "torchvision", "Pillow", "numpy"]
# ///
"""DiffThinker C6: Data scaling on HF Jobs T4. ResNet-18 path-existence classifier."""
import torch, json, os, time, numpy as np
from collections import deque
from torch.utils.data import Dataset, DataLoader
from PIL import Image, ImageDraw

def bfs(grid, start, goal):
    q = deque([start]); v = {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<grid.shape[0] and 0<=nc<grid.shape[1] and grid[nr,nc]==0 and (nr,nc) not in v:
                v.add((nr,nc)); q.append((nr,nc))
    return False

def make_data(n, gs=8):
    data, labels = [], []
    solvable = 0; target = n//2
    while len(data) < n:
        grid = np.zeros((gs,gs), dtype=np.uint8)
        start = (0, np.random.randint(0,gs))
        goal = (gs-1, np.random.randint(0,gs))
        for _ in range(int(gs*gs*np.random.uniform(0.1,0.35))):
            wx,wy = np.random.randint(0,gs,2)
            if (wx,wy)!=start and (wx,wy)!=goal: grid[wx,wy]=1
        hp = bfs(grid, start, goal)
        if hp and solvable >= target: continue
        if not hp and len(data)-solvable >= n-target: continue
        if hp: solvable += 1
        img = Image.new("RGB",(64,64),(255,255,255)); draw=ImageDraw.Draw(img); cw=64//gs
        for r in range(gs):
            for c in range(gs):
                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); labels.append(1 if hp else 0)
    return data, labels

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

device = torch.device("cuda")
gpu = torch.cuda.get_device_name(0)
print(f"GPU: {gpu} | Mem: {torch.cuda.get_device_properties(0).total_memory/1e9:.1f}GB")

from torchvision.models import resnet18, ResNet18_Weights

for n in [20, 50, 100, 200, 400, 800]:
    print(f"\n=== N={n} ===")
    train_imgs, train_lbls = make_data(n)
    test_imgs, test_lbls = make_data(200)
    print(f"  Train: {sum(train_lbls)}/{len(train_lbls)} solvable | Test: {sum(test_lbls)}/200")
    model = resnet18(weights=ResNet18_Weights.DEFAULT).to(device)
    model.fc = torch.nn.Linear(model.fc.in_features, 2)
    model = model.to(device)
    opt = torch.optim.AdamW(model.parameters(), lr=1e-4)
    loader = DataLoader(DS(train_imgs, train_lbls), batch_size=16, shuffle=True)
    t0 = time.time()
    for ep in range(15):
        model.train(); ls=0
        for imgs, lbls in loader:
            imgs, lbls = imgs.to(device), lbls.to(device)
            opt.zero_grad(); out = model(imgs)
            loss = torch.nn.CrossEntropyLoss()(out, lbls)
            loss.backward(); opt.step(); ls+=loss.item()
    model.eval(); correct=0
    with torch.no_grad():
        for i in range(len(test_imgs)):
            img = torch.tensor(np.array(test_imgs[i]).transpose(2,0,1), dtype=torch.float32).unsqueeze(0).to(device)/255.0
            if model(img).argmax(1).item() == test_lbls[i]: correct+=1
    acc = correct/200*100; elapsed = time.time()-t0
    print(f"  Acc: {acc:.1f}% | Time: {elapsed:.1f}s")
    print(json.dumps({"num_train":n,"accuracy":round(acc,1),"time_sec":round(elapsed,1),"gpu":gpu}))