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 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108
| """Train with more data + mild augmentation + bigger model.""" import os, random import numpy as np import torch import torch.nn as nn from PIL import Image, ImageDraw
torch.manual_seed(42) np.random.seed(42) random.seed(42)
DS = 'dataset' OUT = 'model.pt' DEV = 'cuda' if torch.cuda.is_available() else 'cpu'
def get_slot(a, i): x0 = 10 + 42 * i s = a[:, x0:x0+42] if s.shape[1] < 42: s = np.pad(s, ((0, 0), (0, 42 - s.shape[1]))) return s
def load_data(): xs, ys = [], [] with open(os.path.join(DS, 'answers.tsv')) as f: rows = [l.split() for l in f if l.strip()] for fn, ans in rows: a = np.array(Image.open(os.path.join(DS, 'images', fn)).convert('L'), dtype=np.float32) for i, ch in enumerate(ans): xs.append(get_slot(a, i)) ys.append(ord(ch) - 65) X = np.stack(xs) / 255.0 Y = np.array(ys) return X, Y
def aug(x): im = (1.0 - x) * 255.0 img = Image.fromarray(im.astype(np.uint8)) img = img.rotate(random.uniform(-6, 6), fillcolor=0) dx = random.randint(-2, 2); dy = random.randint(-2, 2) img = img.transform(img.size, Image.AFFINE, (1, 0, dx, 0, 1, dy), fillcolor=0) if random.random() < 0.5: d = ImageDraw.Draw(img) for _ in range(random.randint(1, 2)): d.line([(random.randint(0, 41), random.randint(0, 41)), (random.randint(0, 41), random.randint(0, 41))], fill=random.randint(110, 170), width=random.randint(1, 2)) arr = np.array(img, dtype=np.float32) * random.uniform(0.9, 1.1) return 1.0 - arr / 255.0
class Net(nn.Module): def __init__(self): super().__init__() self.net = nn.Sequential( nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Flatten(), nn.Linear(128 * 5 * 5, 256), nn.ReLU(), nn.Dropout(0.4), nn.Linear(256, 26), ) def forward(self, x): return self.net(x)
def main(): X, Y = load_data() print(f'data: {X.shape}, device {DEV}') n_img = len(X) // 5 rng = np.random.RandomState(1) val_img_ids = set(rng.choice(n_img, max(1, n_img // 8), replace=False)) tr_idx = [i for i in range(len(X)) if i // 5 not in val_img_ids] va_idx = [i for i in range(len(X)) if i // 5 in val_img_ids] Xtr, Ytr, Xva, Yva = X[tr_idx], Y[tr_idx], X[va_idx], Y[va_idx] print(f'train {len(Xtr)} val {len(Xva)}')
model = Net().to(DEV) opt = torch.optim.Adam(model.parameters(), lr=2e-3) lossf = nn.CrossEntropyLoss() Xva_t = torch.tensor(Xva, dtype=torch.float32).unsqueeze(1).to(DEV) Yva_t = torch.tensor(Yva, dtype=torch.long).to(DEV)
best_va = 0; best_state = None for epoch in range(40): model.train() perm = np.random.permutation(len(Xtr)) for start in range(0, len(perm), 128): idx = perm[start:start+128] bx = np.stack([aug(Xtr[i]) for i in idx]) xb = torch.tensor(bx, dtype=torch.float32).unsqueeze(1).to(DEV) yb = torch.tensor(Ytr[idx], dtype=torch.long).to(DEV) opt.zero_grad() loss = lossf(model(xb), yb) loss.backward(); opt.step() model.eval() with torch.no_grad(): va_acc = (model(Xva_t).argmax(1) == Yva_t).float().mean().item() print(f'epoch {epoch}: val_acc {va_acc:.4f}') if va_acc > best_va: best_va = va_acc best_state = {k: v.clone() for k, v in model.state_dict().items()} if va_acc > 0.98: break torch.save({'model': best_state}, OUT) print(f'saved {OUT} best_val={best_va:.4f}')
if __name__ == '__main__': main()
|