File size: 2,308 Bytes
963dad2 fe43f4f | 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 | # train every model behind the paper and keep the checkpoint (nothing discarded).
# same configs and seed as the table/ablation runners, so the numbers reproduce.
import os, sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from src.data.perturb_data import load_dataset, MATCH_STRATEGIES
from src.models.encoders import REP_MODES
from src.training.train import TrainConfig, train, save_checkpoint
GPU = int(os.environ.get("PIVOT_GPU", "3"))
ROOT = "models"
def go(name, dataset, data, **cfgkw):
out = os.path.join(ROOT, name)
if os.path.exists(os.path.join(out, "model.pt")):
print("skip", name, flush=True); return
cfg = TrainConfig(dataset=dataset, device_index=GPU, **cfgkw)
model, info = train(cfg, data=data, verbose=False)
save_checkpoint(model, cfg, info, out)
print("saved", name, flush=True)
# main table models: default config (epochs 60), one per dataset and split
splits = {"norman": ["cell", "perturbation", "combination"],
"replogle_k562": ["cell", "perturbation", "gene"]}
for ds, sps in splits.items():
data = load_dataset(ds)
for sp in sps:
go("%s/%s" % (ds, sp), ds, data, split=sp)
# ablation models: norman, perturbation split, epochs 45 (matches run_ablations).
# the "default" run here is the shared base for the inverse-search, guidance-step,
# reward and cost ablations (they reuse one model), plus the components=full,
# representation=gene_op, data-fraction=1.0 and matching=batch rows.
data = load_dataset("norman")
A = "ablations/norman_perturbation"
go("%s/default" % A, "norman", data, split="perturbation", epochs=45)
for nm, comps in [("comp_map", ["map"]), ("comp_map_semi", ["map", "semi"]), ("comp_map_tan", ["map", "tan"])]:
go("%s/%s" % (A, nm), "norman", data, split="perturbation", epochs=45, components=comps)
for rep in REP_MODES:
if rep != "gene_op":
go("%s/rep_%s" % (A, rep), "norman", data, split="perturbation", epochs=45, rep_mode=rep)
for frac in [0.1, 0.25, 0.5, 0.75]:
go("%s/frac_%s" % (A, frac), "norman", data, split="perturbation", epochs=45, train_frac=frac)
for ms in MATCH_STRATEGIES:
if ms != "batch":
go("%s/match_%s" % (A, ms), "norman", data, split="perturbation", epochs=45, match=ms)
print("done", flush=True)
|