| |
| """ |
| hf_train_1kb.py |
| ------------------ |
| HF Jobs version of modal_pipeline_1kb.py. Combines prepare_data() and |
| train_model() into one script, since HF Jobs runs a single command to |
| completion rather than exposing persistent, independently-callable |
| Functions the way Modal does -- there's no direct equivalent of |
| `train_model.remote()` to call repeatedly against a warm environment. |
| |
| All the actual logic (blacklist handling, bedpe parsing with the chr1 |
| fix, span-fit window placement, anchor-weighted masked loss) is copied |
| verbatim from modal_pipeline_1kb.py, which was already tested against |
| synthetic data -- this file only changes HOW it's invoked (plain |
| function calls + argparse, not Modal decorators) and WHERE data lives |
| (a mounted bucket path instead of a Modal Volume path). |
| |
| Expects two mounted volumes (see submit_hf_job.py): |
| /code -- read-only dataset repo containing src/ and configs/ |
| /data -- read-write bucket, plays the role all three Modal volumes |
| (hic-data, hic-cache, hic-runs) played: raw/ downloads, |
| cache/windows_1kb/ per-window .npz cache, runs/ checkpoints. |
| |
| Run (inside the job, via submit_hf_job.py): |
| python hf_train_1kb.py --model bilstm |
| python hf_train_1kb.py --model bilstm --exclude-track RAD21 |
| python hf_train_1kb.py --prepare-only |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import gzip |
| import hashlib |
| import json |
| import math |
| import random |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
| import yaml |
|
|
|
|
| def _cfg(config_dir: Path, config_name: str): |
| with open(config_dir / config_name) as f: |
| return yaml.safe_load(f) |
|
|
|
|
| def _track_fingerprint(track_names: list[str]) -> str: |
| h = hashlib.sha1("|".join(sorted(track_names)).encode()).hexdigest() |
| return h[:8] |
|
|
|
|
| |
|
|
| def load_blacklist(bed_gz_path: Path) -> dict: |
| regions: dict = {} |
| with gzip.open(bed_gz_path, "rt") as fh: |
| for line in fh: |
| if line.startswith("#") or line.startswith("track"): |
| continue |
| p = line.rstrip("\n").split("\t") |
| if len(p) < 3: |
| continue |
| chrom = p[0] if p[0].startswith("chr") else f"chr{p[0]}" |
| try: |
| start, end = int(p[1]), int(p[2]) |
| except ValueError: |
| continue |
| regions.setdefault(chrom, []).append((start, end)) |
| for chrom in regions: |
| regions[chrom].sort() |
| return regions |
|
|
|
|
| def overlaps_blacklist(chrom: str, start: int, end: int, blacklist: dict) -> bool: |
| for (bl_start, bl_end) in blacklist.get(chrom, []): |
| if start < bl_end and end > bl_start: |
| return True |
| return False |
|
|
|
|
| |
|
|
| def parse_encode_bedpe_gz(gz_path) -> list[dict]: |
| loops = [] |
| with gzip.open(gz_path, "rt") as fh: |
| for line in fh: |
| if line.startswith("#"): |
| continue |
| p = line.strip().split("\t") |
| if len(p) < 6: |
| continue |
| c1, c2 = p[0], p[3] |
| try: |
| x1, x2, y1, y2 = int(p[1]), int(p[2]), int(p[4]), int(p[5]) |
| except ValueError: |
| continue |
| if c1 != c2: |
| continue |
| chrom = c1 if c1.startswith("chr") else f"chr{c1}" |
| loops.append({"chrom": chrom, "mid_x": (x1 + x2) // 2, "mid_y": (y1 + y2) // 2}) |
| return loops |
|
|
|
|
| |
|
|
| def place_window_loop_centered(loop, window_size, flank_margin, jitter, rng, sizes, blacklist, |
| max_attempts: int = 5): |
| chrom = loop["chrom"] |
| mid_x, mid_y = loop["mid_x"], loop["mid_y"] |
| span = abs(mid_y - mid_x) |
| max_span = window_size - 2 * flank_margin |
| if span > max_span: |
| return None, "span_too_large" |
| clen = sizes.get(chrom) |
| if clen is None: |
| return None, "unknown_chrom" |
| true_center = (mid_x + mid_y) // 2 |
| for _ in range(max_attempts): |
| center = true_center + rng.randint(-jitter, jitter) |
| start = center - window_size // 2 |
| end = start + window_size |
| if start < 0 or end > clen: |
| continue |
| if not (start + flank_margin <= mid_x <= end - flank_margin): |
| continue |
| if not (start + flank_margin <= mid_y <= end - flank_margin): |
| continue |
| if overlaps_blacklist(chrom, start, end, blacklist): |
| continue |
| return (start, end), None |
| return None, "placement_failed" |
|
|
|
|
| def sample_background_window(chroms, window_size, sizes, blacklist, rng, max_attempts=20): |
| chrom_list = list(chroms) |
| for _ in range(max_attempts): |
| chrom = rng.choice(chrom_list) |
| clen = sizes.get(chrom) |
| if clen is None or clen <= window_size: |
| continue |
| start = rng.randint(0, clen - window_size) |
| end = start + window_size |
| if overlaps_blacklist(chrom, start, end, blacklist): |
| continue |
| return chrom, start, end |
| return None |
|
|
|
|
| |
|
|
| def masked_huber_loss(pred, target, valid_mask, delta: float = 0.5): |
| import torch |
| diff = pred - target |
| abs_diff = diff.abs() |
| quadratic = torch.clamp(abs_diff, max=delta) |
| linear = abs_diff - quadratic |
| per_pixel = (0.5 * quadratic ** 2 + delta * linear) * valid_mask |
| denom = valid_mask.sum().clamp(min=1.0) |
| return per_pixel.sum() / denom |
|
|
|
|
| def masked_mse_loss(pred, target, valid_mask): |
| per_pixel = (pred - target) ** 2 * valid_mask |
| denom = valid_mask.sum().clamp(min=1.0) |
| return per_pixel.sum() / denom |
|
|
|
|
| def build_anchor_weight_map(n_bins: int, anchor_r, anchor_c, boost_radius: int = 5, |
| boost_factor: float = 10.0): |
| w = np.ones((n_bins, n_bins), dtype=np.float32) |
| if anchor_r is not None and anchor_c is not None and anchor_r >= 0 and anchor_c >= 0: |
| r0, r1 = max(0, anchor_r - boost_radius), min(n_bins, anchor_r + boost_radius + 1) |
| c0, c1 = max(0, anchor_c - boost_radius), min(n_bins, anchor_c + boost_radius + 1) |
| w[r0:r1, c0:c1] = boost_factor |
| return w |
|
|
|
|
| def masked_weighted_huber_loss(pred, target, valid_mask, weight_map, delta: float = 0.5): |
| import torch |
| diff = pred - target |
| abs_diff = diff.abs() |
| quadratic = torch.clamp(abs_diff, max=delta) |
| linear = abs_diff - quadratic |
| per_pixel = (0.5 * quadratic ** 2 + delta * linear) * valid_mask * weight_map |
| denom = (valid_mask * weight_map).sum().clamp(min=1.0) |
| return per_pixel.sum() / denom |
|
|
|
|
| |
|
|
| def resolve_file(rel_path: str, shared_raw, local_raw: Path, url, genomic_io) -> Path: |
| """Returns a usable local path for `rel_path`, preferring a file that |
| already exists over downloading it. |
| |
| Lookup order: |
| 1. shared_raw/rel_path -- the read-only data bucket (e.g. the existing |
| 3d-genomics-prediction bucket, which already holds every bigWig and |
| both .hic files) |
| 2. local_raw/rel_path -- downloaded by an earlier run |
| 3. download from `url` into local_raw/rel_path |
| |
| This exists because the first HF Jobs run re-downloaded ~90GB from |
| ENCODE into a fresh bucket while identical files sat unused in an |
| existing one. Checking before fetching turns that into a no-op. |
| Raises if missing everywhere AND no url is given, rather than returning |
| a path that doesn't exist and failing later inside pyBigWig/hicstraw |
| with a much less obvious message. |
| """ |
| if shared_raw is not None: |
| candidate = Path(shared_raw) / rel_path |
| if candidate.exists() and candidate.stat().st_size > 0: |
| print(f" [shared] {rel_path}") |
| return candidate |
| candidate = local_raw / rel_path |
| if candidate.exists() and candidate.stat().st_size > 0: |
| print(f" [cache] {rel_path}") |
| return candidate |
| if not url: |
| raise FileNotFoundError( |
| f"'{rel_path}' not found in shared raw dir ({shared_raw}) or local raw dir " |
| f"({local_raw}), and no download url was configured for it.") |
| genomic_io.download_file(url, candidate) |
| return candidate |
|
|
|
|
| def prepare_data(cfg: dict, data_dir: Path, genomic_io, genome_constants, |
| shared_raw=None) -> None: |
| exp = cfg["experiment"] |
| raw = data_dir / "raw" |
| raw.mkdir(parents=True, exist_ok=True) |
| manifest_path = data_dir / "loop_manifest_1kb.csv" |
|
|
| track_name_sets = {cl: sorted(cl_cfg["tracks"].keys()) for cl, cl_cfg in cfg["cell_lines"].items()} |
| reference_tracks = next(iter(track_name_sets.values())) |
| for cl, names in track_name_sets.items(): |
| if names != reference_tracks: |
| raise ValueError(f"Track set mismatch: '{cl}' has {names}, expected {reference_tracks}.") |
| track_fp = _track_fingerprint(reference_tracks) |
|
|
| fingerprint_path = data_dir / "manifest_fingerprint_1kb.txt" |
| needed_fp = (f"{sorted(cfg['cell_lines'].keys())}|{track_fp}|" |
| f"{exp['bin_size']}|{exp['window_size']}|{exp['flank_margin']}") |
|
|
| blacklist = {} |
| if cfg.get("blacklist_filename"): |
| bl_path = resolve_file(cfg["blacklist_filename"], shared_raw, raw, |
| cfg.get("blacklist_url"), genomic_io) |
| blacklist = load_blacklist(bl_path) |
| n_regions = sum(len(v) for v in blacklist.values()) |
| print(f"[prepare] Blacklist: {n_regions} regions across {len(blacklist)} chroms.") |
|
|
| for cl_name, cl_cfg in cfg["cell_lines"].items(): |
| print(f"\n[prepare] {cl_name}: resolving {len(cl_cfg['tracks'])} tracks, loop calls, " |
| f"and the intact .hic file...") |
| for track_name, entry in cl_cfg["tracks"].items(): |
| resolve_file(entry["path"], shared_raw, raw, entry.get("url"), genomic_io) |
| resolve_file(cl_cfg["loop_filename"], shared_raw, raw, cl_cfg.get("loop_url"), genomic_io) |
| |
| |
| |
| |
| resolve_file(cl_cfg.get("hic_path", f"{cl_name}.hic"), shared_raw, raw, |
| cl_cfg.get("hic_url"), genomic_io) |
| print(f"[prepare] {cl_name}: done.") |
|
|
| if manifest_path.exists(): |
| old_fp = fingerprint_path.read_text().strip() if fingerprint_path.exists() else None |
| if old_fp == needed_fp: |
| print(f"\n[prepare] Manifest exists for {needed_fp}. Downloads verified above; " |
| f"skipping window rebuild.") |
| return |
| print(f"\n[prepare] Config changed ({old_fp} -> {needed_fp}). Rebuilding manifest.") |
| manifest_path.unlink() |
|
|
| sizes = genome_constants.chrom_sizes_hg38() |
| bin_size = exp["bin_size"]; window_size = exp["window_size"] |
| jitter = exp["jitter"]; flank_margin = exp["flank_margin"] |
| n_bins = window_size // bin_size |
| train_chroms = set(exp["train_chroms"]); test_chroms = set(exp["test_chroms"]) |
| rng = random.Random(exp["seed"]) |
|
|
| rows = [] |
| n_span_rejected = 0; n_placement_failed = 0 |
| for cl_name, cl_cfg in cfg["cell_lines"].items(): |
| loop_file = resolve_file(cl_cfg["loop_filename"], shared_raw, raw, |
| cl_cfg.get("loop_url"), genomic_io) |
| loops = parse_encode_bedpe_gz(loop_file) |
| print(f"\n[prepare] {cl_name}: {len(loops)} intra-chromosomal loops parsed.") |
|
|
| for split, chroms, n_target in [ |
| ("train", train_chroms, exp["train_loops_per_cell_line"]), |
| ("test", test_chroms, exp["test_loops_per_cell_line"]), |
| ]: |
| eligible = [l for l in loops if l["chrom"] in chroms] |
| rng.shuffle(eligible) |
| placed = 0 |
| for loop in eligible: |
| if placed >= n_target: |
| break |
| result, reason = place_window_loop_centered( |
| loop, window_size, flank_margin, jitter, rng, sizes, blacklist) |
| if result is None: |
| if reason == "span_too_large": |
| n_span_rejected += 1 |
| else: |
| n_placement_failed += 1 |
| continue |
| start, end = result |
| rows.append({ |
| "split": split, "cell_line": cl_name, "window_type": "loop", |
| "chrom": loop["chrom"], "start": start, "end": end, |
| "bin_size": bin_size, "n_bins": n_bins, |
| "loop_mid_x": loop["mid_x"], "loop_mid_y": loop["mid_y"], |
| "loop_span_bp": abs(loop["mid_y"] - loop["mid_x"]), |
| "track_fingerprint": track_fp, |
| }) |
| placed += 1 |
| status = "WARNING: only" if placed < n_target else " " |
| print(f" {status} {placed}/{n_target} loop windows: {cl_name}/{split}") |
|
|
| n_bg_target = exp.get("background_windows_per_cell_line", 0) |
| placed_bg = 0 |
| for _ in range(n_bg_target * 3): |
| if placed_bg >= n_bg_target: |
| break |
| result = sample_background_window(train_chroms, window_size, sizes, blacklist, rng) |
| if result is None: |
| continue |
| chrom, start, end = result |
| rows.append({ |
| "split": "train", "cell_line": cl_name, "window_type": "background", |
| "chrom": chrom, "start": start, "end": end, |
| "bin_size": bin_size, "n_bins": n_bins, |
| "loop_mid_x": "", "loop_mid_y": "", "loop_span_bp": "", |
| "track_fingerprint": track_fp, |
| }) |
| placed_bg += 1 |
| print(f" {placed_bg}/{n_bg_target} background windows: {cl_name}/train") |
|
|
| print(f"\n[prepare] Rejected: {n_span_rejected} loops (span too large), " |
| f"{n_placement_failed} (blacklist/bounds after retries).") |
|
|
| fieldnames = ["split", "cell_line", "window_type", "chrom", "start", "end", |
| "bin_size", "n_bins", "loop_mid_x", "loop_mid_y", "loop_span_bp", |
| "track_fingerprint"] |
| with open(manifest_path, "w", newline="") as fh: |
| w = csv.DictWriter(fh, fieldnames=fieldnames) |
| w.writeheader(); w.writerows(rows) |
| fingerprint_path.write_text(needed_fp) |
| print(f"\n[prepare] Done: {len(rows)} windows -> {manifest_path}") |
|
|
|
|
| |
|
|
| def train_model(model_name: str, cfg: dict, data_dir: Path, src_dir: Path, |
| exclude_track: str = "", shared_raw=None) -> dict: |
| import torch |
| from torch.utils.data import DataLoader, Dataset |
| sys.path.insert(0, str(src_dir)) |
| from genomic_io import read_bigwig, _get_hicfile |
| from oe_extraction import get_oe_matrix, MAX_OE |
|
|
| tr_cfg = cfg["training"] |
| raw = data_dir / "raw" |
|
|
| def _resolve_existing(rel_path: str) -> str: |
| """Same shared-then-local lookup as resolve_file(), minus the |
| download fallback -- by training time prepare_data() has already |
| ensured every file exists somewhere, so a miss here is a bug worth |
| surfacing immediately rather than silently re-fetching 40GB mid-run.""" |
| if shared_raw is not None: |
| cand = Path(shared_raw) / rel_path |
| if cand.exists(): |
| return str(cand) |
| return str(raw / rel_path) |
|
|
| train_seed = int(tr_cfg.get("seed", cfg.get("experiment", {}).get("seed", 42))) |
| torch.manual_seed(train_seed) |
| torch.cuda.manual_seed_all(train_seed) |
| print(f"[train] seed={train_seed}") |
|
|
| track_names = sorted(next(iter(cfg["cell_lines"].values()))["tracks"].keys()) |
| if exclude_track: |
| if exclude_track not in track_names: |
| raise ValueError(f"exclude_track='{exclude_track}' not in {track_names}") |
| track_names = [t for t in track_names if t != exclude_track] |
| print(f"[train] LEAVE-ONE-TRACK-OUT: excluding '{exclude_track}'") |
| n_tracks = len(track_names) |
| track_fp = _track_fingerprint(track_names) |
| max_oe = float(tr_cfg.get("max_oe", MAX_OE)) |
| print(f"[train] {n_tracks} tracks: {track_names}") |
|
|
| class LoopDataset(Dataset): |
| def __init__(self, manifest_path, split, normalizer=None, cache_dir=None, clip_pct=99.0): |
| full = pd.read_csv(manifest_path) |
| self.manifest = full[full["split"] == split].reset_index(drop=True) |
| self.normalizer = normalizer; self.clip_pct = clip_pct |
| self.cache_root = Path(cache_dir) if cache_dir else None |
| if self.cache_root: |
| self.cache_root.mkdir(parents=True, exist_ok=True) |
|
|
| def __len__(self): |
| return len(self.manifest) |
|
|
| def _paths(self, cl): |
| cl_tracks = cfg["cell_lines"][cl]["tracks"] |
| return {name: _resolve_existing(cl_tracks[name]["path"]) for name in track_names} | \ |
| {"hic": _resolve_existing(cfg["cell_lines"][cl].get("hic_path", f"{cl}.hic"))} |
|
|
| def _cache(self, row): |
| if self.cache_root is None: |
| return None |
| bs = int(row["bin_size"]) |
| return self.cache_root / (f"{row['cell_line']}_{row['chrom']}_{row['start']}_" |
| f"{row['end']}_log2oe_v2_{bs}_{track_fp}.npz") |
|
|
| def get_tracks(self, idx): |
| row = self.manifest.iloc[idx] |
| p = self._paths(row["cell_line"]); n = int(row["n_bins"]) |
| ch, s, e = row["chrom"], int(row["start"]), int(row["end"]) |
| arrays = [read_bigwig(p[name], ch, s, e, n, self.clip_pct) for name in track_names] |
| return np.stack(arrays, axis=0) |
|
|
| def __getitem__(self, idx): |
| row = self.manifest.iloc[idx]; cl = row["cell_line"] |
| ch, s, e = row["chrom"], int(row["start"]), int(row["end"]) |
| bs = int(row["bin_size"]); cf = self._cache(row) |
| if cf is not None and cf.exists(): |
| data = np.load(cf) |
| tracks, log2_oe, valid_mask = data["tracks"], data["log2_oe"], data["valid_mask"] |
| else: |
| tracks = self.get_tracks(idx) |
| hf = _get_hicfile(self._paths(cl)["hic"]) |
| _, log2_oe, valid_mask = get_oe_matrix(hf, ch, s, e, bs, max_oe=max_oe) |
| if cf is not None: |
| np.savez_compressed(cf, tracks=tracks.astype(np.float32), |
| log2_oe=log2_oe, valid_mask=valid_mask) |
| if self.normalizer is not None: |
| s_n = self.normalizer[cl] |
| tracks = (tracks - s_n["mean"][:, None]) / s_n["std"][:, None] |
| if row["window_type"] == "loop" and pd.notna(row.get("loop_mid_x")): |
| anchor_r = (int(row["loop_mid_x"]) - s) // bs |
| anchor_c = (int(row["loop_mid_y"]) - s) // bs |
| else: |
| anchor_r = anchor_c = -1 |
| return {"tracks": torch.from_numpy(tracks).float(), |
| "target": torch.from_numpy(log2_oe).float(), |
| "valid_mask": torch.from_numpy(valid_mask.astype(np.float32)), |
| "anchor_r": anchor_r, "anchor_c": anchor_c, "cell_line": cl} |
|
|
| manifest = data_dir / "loop_manifest_1kb.csv" |
| run_dir = data_dir / "runs" / f"1kb_{model_name}_{track_fp}" |
| run_dir.mkdir(parents=True, exist_ok=True) |
| cache_dir = data_dir / "cache" / "windows_1kb" |
|
|
| norm_path = data_dir / "runs" / f"normalizer_1kb_{track_fp}.json" |
| config_cls = set(cfg["cell_lines"].keys()) |
| need_refit = True |
| if norm_path.exists(): |
| with open(norm_path) as f: |
| normalizer = json.load(f) |
| if set(normalizer.keys()) == config_cls: |
| need_refit = False |
| print(f"[train] Loaded normalizer for {config_cls}.") |
|
|
| if need_refit: |
| train_df = pd.read_csv(manifest) |
| train_df = train_df[train_df["split"] == "train"].reset_index(drop=True) |
| raw_ds = LoopDataset(manifest, "train", normalizer=None, cache_dir=cache_dir) |
| normalizer = {} |
| for cl in train_df["cell_line"].unique(): |
| cl_pos = train_df[train_df["cell_line"] == cl].index.tolist() |
| n_s = min(80, len(cl_pos)) |
| sample_idx = np.linspace(0, len(cl_pos) - 1, n_s).astype(int) |
| sums = np.zeros(n_tracks); sq_sums = np.zeros(n_tracks); count = 0 |
| for si in sample_idx: |
| t = raw_ds.get_tracks(cl_pos[si]) |
| sums += t.sum(axis=1); sq_sums += (t ** 2).sum(axis=1); count += t.shape[1] |
| mean = sums / count |
| std = np.sqrt(np.clip(sq_sums / count - mean ** 2, 1e-8, None)) |
| normalizer[cl] = {"mean": mean.tolist(), "std": std.tolist()} |
| print(f" {cl}: " + " ".join(f"{n}={m:.3f}" for n, m in zip(track_names, mean))) |
| run_dir.parent.mkdir(parents=True, exist_ok=True) |
| with open(norm_path, "w") as f: |
| json.dump(normalizer, f, indent=2) |
|
|
| norm_np = {cl: {"mean": np.array(v["mean"], dtype=np.float32), |
| "std": np.array(v["std"], dtype=np.float32)} for cl, v in normalizer.items()} |
|
|
| train_ds = LoopDataset(manifest, "train", normalizer=norm_np, cache_dir=cache_dir) |
| test_ds = LoopDataset(manifest, "test", normalizer=norm_np, cache_dir=cache_dir) |
| print(f"[train] {len(train_ds)} train, {len(test_ds)} test windows") |
|
|
| bs = tr_cfg["batch_size"] |
| train_loader = DataLoader(train_ds, batch_size=bs, shuffle=True, num_workers=0, pin_memory=True) |
| test_loader = DataLoader(test_ds, batch_size=bs, shuffle=False, num_workers=0, pin_memory=True) |
| n_bins = int(train_ds.manifest["n_bins"].iloc[0]) |
|
|
| device = "cuda" |
| if model_name == "cnn": |
| from model_cnn import CNNHiCModel |
| model = CNNHiCModel(n_tracks=n_tracks, n_bins=n_bins).to(device) |
| wd = float(tr_cfg["weight_decay"]); clip_norm = None |
| elif model_name == "unet": |
| from model_unet import UNetHiCModel |
| model = UNetHiCModel(n_tracks=n_tracks, n_bins=n_bins).to(device) |
| wd = float(tr_cfg["weight_decay"]); clip_norm = None |
| elif model_name == "transformer": |
| from model_transformer import TransformerHiCModel |
| dropout = float(tr_cfg.get("transformer_dropout", 0.3)) |
| model = TransformerHiCModel(n_tracks=n_tracks, n_bins=n_bins, dropout=dropout).to(device) |
| wd = float(tr_cfg.get("transformer_weight_decay", 1e-3)) |
| clip_norm = float(tr_cfg.get("transformer_grad_clip_norm", 0.5)) |
| elif model_name == "bilstm": |
| from model_bilstm import BiLSTMHiCModel |
| dropout = float(tr_cfg.get("bilstm_dropout", 0.2)) |
| model = BiLSTMHiCModel(n_tracks=n_tracks, n_bins=n_bins, dropout=dropout).to(device) |
| wd = float(tr_cfg.get("bilstm_weight_decay", 1e-5)) |
| clip_norm = float(tr_cfg.get("bilstm_grad_clip_norm", 1.0)) |
| else: |
| raise ValueError(f"Unknown model: {model_name}") |
|
|
| n_params = sum(p.numel() for p in model.parameters()) |
| print(f"[train] {model_name}: {n_params:,} params") |
|
|
| epochs = tr_cfg["epochs"][model_name] |
| patience = tr_cfg.get("early_stopping_patience", 10) |
| warmup = tr_cfg.get("transformer_warmup_steps", 0) if model_name == "transformer" else 0 |
| total_steps = len(train_loader) * epochs |
| optimizer = torch.optim.Adam(model.parameters(), lr=float(tr_cfg["lr"]), weight_decay=wd) |
|
|
| def lr_lambda(step): |
| if warmup <= 0: |
| return 1.0 |
| if step < warmup: |
| return (step + 1) / warmup |
| progress = min(1.0, (step - warmup) / max(1, total_steps - warmup)) |
| return max(0.05, 0.5 * (1 + math.cos(math.pi * progress))) |
|
|
| scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) |
|
|
| log_path = run_dir / "training_log.csv" |
| ckpt_path = run_dir / "best_model.pt" |
| best_test_loss = float("inf"); epochs_no_improve = 0; stopped_at = None |
| huber_delta = float(tr_cfg.get("huber_delta", 0.5)) |
| anchor_boost_radius = int(tr_cfg.get(f"{model_name}_anchor_boost_radius", |
| tr_cfg.get("anchor_boost_radius", 5))) |
| anchor_boost_factor = float(tr_cfg.get(f"{model_name}_anchor_boost_factor", |
| tr_cfg.get("anchor_boost_factor", 10.0))) |
| print(f"[train] anchor-weighted loss: radius={anchor_boost_radius} factor={anchor_boost_factor}") |
|
|
| fields = ["epoch", "train_huber", "train_huber_weighted", "train_mse", "test_huber", |
| "test_mse", "epoch_seconds", "lr"] |
| with open(log_path, "w", newline="") as log_fh: |
| writer = csv.DictWriter(log_fh, fieldnames=fields) |
| writer.writeheader() |
| for epoch in range(1, epochs + 1): |
| t0 = time.time() |
| model.train() |
| tr_hub = tr_hub_weighted = tr_mse = n_tr = 0; last_lr = None |
| for batch in train_loader: |
| tracks = batch["tracks"].to(device, non_blocking=True) |
| target = batch["target"].to(device, non_blocking=True) |
| valid = batch["valid_mask"].to(device, non_blocking=True) |
| anchors_r = batch["anchor_r"].numpy(); anchors_c = batch["anchor_c"].numpy() |
| weight_maps = np.stack([ |
| build_anchor_weight_map(n_bins, r if r >= 0 else None, c if c >= 0 else None, |
| boost_radius=anchor_boost_radius, |
| boost_factor=anchor_boost_factor) |
| for r, c in zip(anchors_r, anchors_c)]) |
| weight = torch.from_numpy(weight_maps).float().to(device, non_blocking=True) |
|
|
| optimizer.zero_grad() |
| pred = model(tracks) |
| loss = masked_weighted_huber_loss(pred, target, valid, weight, delta=huber_delta) |
| loss.backward() |
| if clip_norm: |
| torch.nn.utils.clip_grad_norm_(model.parameters(), clip_norm) |
| optimizer.step(); scheduler.step() |
| last_lr = scheduler.get_last_lr()[0] |
| tr_hub_weighted += loss.item() |
| tr_hub += masked_huber_loss(pred, target, valid, delta=huber_delta).item() |
| tr_mse += masked_mse_loss(pred, target, valid).item() |
| n_tr += 1 |
| tr_hub /= max(1, n_tr); tr_hub_weighted /= max(1, n_tr); tr_mse /= max(1, n_tr) |
|
|
| model.eval(); te_hub = te_mse = n_te = 0 |
| with torch.no_grad(): |
| for batch in test_loader: |
| tracks = batch["tracks"].to(device, non_blocking=True) |
| target = batch["target"].to(device, non_blocking=True) |
| valid = batch["valid_mask"].to(device, non_blocking=True) |
| pred = model(tracks) |
| te_hub += masked_huber_loss(pred, target, valid, delta=huber_delta).item() |
| te_mse += masked_mse_loss(pred, target, valid).item() |
| n_te += 1 |
| te_hub /= max(1, n_te); te_mse /= max(1, n_te) |
|
|
| dt = time.time() - t0 |
| print(f" epoch {epoch:3d}/{epochs} tr_hub={tr_hub:.4f} tr_hub_w={tr_hub_weighted:.4f} " |
| f"tr_mse={tr_mse:.4f} te_hub={te_hub:.4f} te_mse={te_mse:.4f} " |
| f"lr={last_lr:.2e} ({dt:.1f}s)") |
| writer.writerow({"epoch": epoch, "train_huber": tr_hub, |
| "train_huber_weighted": tr_hub_weighted, "train_mse": tr_mse, |
| "test_huber": te_hub, "test_mse": te_mse, |
| "epoch_seconds": dt, "lr": last_lr}) |
| log_fh.flush() |
|
|
| if te_hub < best_test_loss: |
| best_test_loss = te_hub; epochs_no_improve = 0 |
| torch.save({"model_name": model_name, "n_bins": n_bins, "n_tracks": n_tracks, |
| "track_names": track_names, "epoch": epoch, |
| "model_state_dict": model.state_dict(), |
| "test_huber": te_hub, "test_mse": te_mse}, ckpt_path) |
| else: |
| epochs_no_improve += 1 |
| if epochs_no_improve >= patience: |
| stopped_at = epoch |
| print(f"[train] Early stop epoch {epoch}") |
| break |
|
|
| return {"model": model_name, "n_tracks": n_tracks, "best_test_huber": best_test_loss, |
| "best_epoch": int(torch.load(ckpt_path, map_location="cpu", weights_only=False)["epoch"]), |
| "stopped_early_at": stopped_at} |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser() |
| p.add_argument("--model", default="bilstm", choices=["cnn", "unet", "transformer", "bilstm"]) |
| p.add_argument("--config-name", default="encode_sources_1kb.yaml") |
| p.add_argument("--exclude-track", default="") |
| p.add_argument("--prepare-only", action="store_true") |
| p.add_argument("--code-dir", default="/code") |
| p.add_argument("--data-dir", default="/data") |
| p.add_argument("--shared-raw-dir", default="", |
| help="Read-only mount holding the already-uploaded raw data " |
| "(e.g. /shared for the 3d-genomics-prediction bucket). Files found " |
| "here are used as-is; anything missing falls back to downloading " |
| "into --data-dir/raw. Empty = download everything, the old behavior.") |
| args = p.parse_args() |
|
|
| code_dir = Path(args.code_dir) |
| data_dir = Path(args.data_dir) |
| shared_raw = Path(args.shared_raw_dir) if args.shared_raw_dir else None |
| src_dir = code_dir / "src" |
| sys.path.insert(0, str(src_dir)) |
| import genomic_io |
| import genome_constants |
|
|
| cfg = _cfg(code_dir / "configs", args.config_name) |
|
|
| print("=" * 65) |
| print(f"HepG2/HCT116 1kb loop prediction (HF Jobs) -- config: {args.config_name}") |
| if shared_raw: |
| print(f"shared read-only raw data: {shared_raw}") |
| print("=" * 65) |
| print("\n[Step 1] Preparing data...") |
| prepare_data(cfg, data_dir, genomic_io, genome_constants, shared_raw=shared_raw) |
| print("[Step 1] Done.") |
| if args.prepare_only: |
| return |
|
|
| print(f"\n[Step 2] Training {args.model}" |
| + (f" (excluding {args.exclude_track})" if args.exclude_track else "") + "...") |
| r = train_model(args.model, cfg, data_dir, src_dir, args.exclude_track, shared_raw=shared_raw) |
| s = f" (early@{r['stopped_early_at']})" if r["stopped_early_at"] else "" |
| print(f"\n<- {args.model.upper()}{s}: best_test_huber={r['best_test_huber']:.4f} " |
| f"epoch={r['best_epoch']}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
