# SPDX-License-Identifier: MIT """YS-EN-0003 ยท train and test an image model on encoded light curves. One run = one (encoding, backbone, init, resample, training-size) combination. Results are appended to results/image_runs.jsonl; predictions to results/preds/. The recipe is fixed in advance and identical for every run (no tuning on the test set): AdamW, 1,000 optimisation steps of batch 32 (about 32 epochs at the full 1,000 training curves), cosine schedule with 50 warm-up steps, label smoothing 0.1, weight decay 0.05, no augmentation, 224 x 224 inputs. Learning rate 2e-4 for pretrained weights, 1e-3 from scratch. Examples: python train_image.py --encoding stack --backbone resnet18 --pretrained --resample 0 python train_image.py --encoding gasf --backbone dinov2_s --probe --resample 0 python train_image.py --original --resample 0 # re-creation of our 2023 sketch """ from __future__ import annotations import argparse import json import math import time from pathlib import Path import numpy as np import timm import torch import torch.nn.functional as F import lcdata HERE = Path(__file__).resolve().parent OUT = HERE / "results" DEV = "cuda" if torch.cuda.is_available() else "cpu" BACKBONES = { # name -> (timm id, kwargs) "resnet18": ("resnet18.tv_in1k", {}), "resnet50": ("resnet50.tv2_in1k", {}), "convnext_t": ("convnext_tiny.fb_in1k", {}), "efficientnet_b0": ("efficientnet_b0.ra_in1k", {}), "vit_s16": ("vit_small_patch16_224.augreg_in21k_ft_in1k", {}), "dinov2_s": ("vit_small_patch14_dinov2.lvd142m", {"img_size": 224}), } MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) def to_input(img: torch.Tensor) -> torch.Tensor: """[-1, 1] encoding -> ImageNet-normalised tensor.""" x = (img + 1) / 2 return (x - MEAN.to(x.device)) / STD.to(x.device) def rainbow(img: torch.Tensor) -> torch.Tensor: """Map a 1-channel [-1, 1] field through matplotlib's 'rainbow' colour map, as the 2023 sketch's exported PNGs did.""" from matplotlib import colormaps lut = torch.tensor(colormaps["rainbow"](np.linspace(0, 1, 256))[:, :3], dtype=torch.float32, device=img.device) i = ((img[:, 0] + 1) / 2 * 255).round().long().clamp(0, 255) return lut[i].permute(0, 3, 1, 2) * 2 - 1 def make_model(args): if args.original: from timm.models.vision_transformer import VisionTransformer # vit_pytorch.ViT(image_size=200, patch_size=50, dim=1024, depth=6, heads=8, # mlp_dim=2048, dropout=0.1, emb_dropout=0.1) re-created in timm return VisionTransformer(img_size=200, patch_size=50, num_classes=3, embed_dim=1024, depth=6, num_heads=8, mlp_ratio=2.0, drop_rate=0.1, pos_drop_rate=0.1) name, kw = BACKBONES[args.backbone] return timm.create_model(name, pretrained=args.pretrained, num_classes=0 if args.probe else 3, **kw) def batches(n, bs, rng): while True: p = rng.permutation(n) for i in range(0, n - bs + 1 if n >= bs else 1, bs): yield p[i:i + bs] @torch.no_grad() def forward_all(model, X, args, bs=64): model.eval() outs = [] for i in range(0, len(X), bs): outs.append(model(prep(X[i:i + bs], args)).float().cpu()) return torch.cat(outs) def prep(xb, args): if args.original: img = rainbow(lcdata.encode(xb, "gasf", 200)[:, :1]) return to_input(img) return to_input(lcdata.encode(xb, args.encoding, args.size)) def run_name(args): if args.original: return "original-vit-rainbow-gasf" return f"{args.encoding}-{args.backbone}-{'probe' if args.probe else ('pt' if args.pretrained else 'scratch')}" def lcdata_ntrain(args): """Size of the stratified subset subsample() returns (rounding per class: 50 -> 51).""" _, y, _, _ = (lcdata.load_ogle if args.dataset == "ogle" else lcdata.load)(args.resample) return len(lcdata.subsample(y, args.ntrain, seed=args.resample)) def main(): ap = argparse.ArgumentParser() ap.add_argument("--encoding", default="stack", choices=["gasf", "gadf", "mtf", "rp", "line", "stack"]) ap.add_argument("--backbone", default="resnet18", choices=list(BACKBONES)) ap.add_argument("--pretrained", action="store_true") ap.add_argument("--probe", action="store_true", help="frozen features + logistic regression") ap.add_argument("--original", action="store_true", help="re-create the 2023 GAF + ViT sketch") ap.add_argument("--resample", type=int, default=0) ap.add_argument("--ntrain", type=int, default=0, help="stratified training subset (0 = all)") ap.add_argument("--steps", type=int, default=1000) ap.add_argument("--bs", type=int, default=32) ap.add_argument("--size", type=int, default=224) ap.add_argument("--tag", default="") ap.add_argument("--dataset", default="slc", choices=["slc", "ogle"]) ap.add_argument("--skip-done", action="store_true", help="exit if this exact run is already in image_runs.jsonl") args = ap.parse_args() if args.probe: args.pretrained = True # a probe of random features is not interesting if args.skip_done and (OUT / "image_runs.jsonl").exists(): key = run_name(args) for line in (OUT / "image_runs.jsonl").read_text().splitlines(): d = json.loads(line) if (d.get("dataset", "slc") == args.dataset and d["run"] == key and d["resample"] == args.resample and (not args.ntrain or d["ntrain"] == lcdata_ntrain(args)) and d.get("tag", "") == args.tag): print("already done:", key, args.resample, args.ntrain) return torch.manual_seed(args.resample) rng = np.random.default_rng(args.resample) X, y, Xt, yt = (lcdata.load_ogle if args.dataset == "ogle" else lcdata.load)(args.resample) if args.ntrain: idx = lcdata.subsample(y, args.ntrain, seed=args.resample) X, y = X[idx], y[idx] Xg = torch.tensor(X, device=DEV) yg = torch.tensor(y, device=DEV) Xtg = torch.tensor(Xt, device=DEV) model = make_model(args).to(DEV) t0 = time.time() if args.probe: from sklearn.linear_model import LogisticRegressionCV from sklearn.preprocessing import StandardScaler ftr = forward_all(model, Xg, args).numpy() fte = forward_all(model, Xtg, args).numpy() sc = StandardScaler().fit(ftr) clf = LogisticRegressionCV(Cs=10, cv=5, max_iter=3000, n_jobs=-1).fit(sc.transform(ftr), y) prob = clf.predict_proba(sc.transform(fte)) else: if args.original: # the notebook: fit_one_cycle(12, 2.09e-5), batch 8, Adam (fastai defaults) bs, steps = 8, 12 * math.ceil(len(X) / 8) opt = torch.optim.AdamW(model.parameters(), lr=2.09e-5, weight_decay=0.01) sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=2.09e-5, total_steps=steps) smooth = 0.0 else: bs, steps = args.bs, args.steps lr = 2e-4 if args.pretrained else 1e-3 opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.05) warm = 50 sched = torch.optim.lr_scheduler.LambdaLR( opt, lambda s: min(1, (s + 1) / warm) * 0.5 * (1 + math.cos(math.pi * min(s, steps) / steps))) smooth = 0.1 model.train() it = batches(len(X), min(bs, len(X)), rng) for step in range(steps): b = torch.tensor(next(it), device=DEV) loss = F.cross_entropy(model(prep(Xg[b], args)), yg[b], label_smoothing=smooth) opt.zero_grad(set_to_none=True) loss.backward() opt.step() sched.step() if step % 200 == 0: print(f"step {step:5d} loss {loss.item():.4f} {time.time() - t0:6.0f}s", flush=True) prob = F.softmax(forward_all(model, Xtg, args), 1).numpy() pred = prob.argmax(1) acc = float((pred == yt).mean()) cm = np.zeros((3, 3), int) for a, b in zip(yt, pred): cm[a, b] += 1 name = run_name(args) rec = { "run": name, "encoding": "gasf-rainbow" if args.original else args.encoding, "backbone": "vit-original" if args.original else args.backbone, "init": "scratch" if (args.original or not args.pretrained) else "imagenet", "mode": "probe" if args.probe else "finetune", "resample": args.resample, "ntrain": len(X), "acc": acc, "confusion": cm.tolist(), "seconds": round(time.time() - t0, 1), "tag": args.tag, "peak_rss_mb": round(__import__("resource").getrusage(__import__("resource").RUSAGE_SELF).ru_maxrss / 1024), } if args.dataset != "slc": rec["dataset"] = args.dataset name = f"{args.dataset}-{name}" OUT.mkdir(exist_ok=True) (OUT / "preds").mkdir(exist_ok=True) np.savez_compressed(OUT / "preds" / f"{name}-r{args.resample}-n{len(X)}{args.tag}.npz", prob=prob.astype(np.float16), y=yt) with open(OUT / "image_runs.jsonl", "a") as f: f.write(json.dumps(rec) + "\n") print(json.dumps(rec)) if __name__ == "__main__": main()