"""Pretrain a TimesNet + ConvNeXtV2 autoencoder on training-set exposure.""" from __future__ import annotations import argparse import json import math from pathlib import Path import numpy as np import torch from torch.optim import AdamW from torch.utils.data import DataLoader, Dataset from tqdm import tqdm from backbones import TimesNetExposureAutoencoder from dataset import ExposureCache from train_util import ( configure_torch_for_training, create_unique_run_dir, load_eid_file, resolve_device, set_seed, setup_logging, ) class ExposureWindowDataset(Dataset): def __init__(self, cache: ExposureCache, row_indices: np.ndarray): self.cache = cache self.row_indices = np.asarray(row_indices, dtype=np.int64) def __len__(self) -> int: return len(self.row_indices) def __getitem__(self, index: int) -> dict[str, torch.Tensor]: row = int(self.row_indices[index]) return { "daily": torch.from_numpy( np.array(self.cache.daily[row], dtype=np.float32, copy=True) ), "monthly": torch.from_numpy( np.array(self.cache.monthly[row], dtype=np.float32, copy=True) ), } def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Pretrain a TimesNet + ConvNeXtV2 exposure autoencoder" ) parser.add_argument("--exposure_cache_dir", required=True) parser.add_argument("--train_eid_file", default="ukb_train_eid.csv") parser.add_argument("--runs_root", default="runs") parser.add_argument("--seed", type=int, default=42) parser.add_argument("--val_fraction", type=float, default=0.05) parser.add_argument("--n_embd", type=int, default=120) parser.add_argument("--d_model", type=int, default=None) parser.add_argument("--n_layers", type=int, default=2) parser.add_argument("--top_k", type=int, default=3) parser.add_argument("--n_convnext_blocks", type=int, default=2) parser.add_argument("--conv_kernel_size", type=int, default=7) parser.add_argument("--mlp_ratio", type=float, default=4.0) parser.add_argument("--dropout", type=float, default=0.0) parser.add_argument("--mask_ratio", type=float, default=0.25) parser.add_argument("--batch_size", type=int, default=16) parser.add_argument("--base_lr", type=float, default=3e-4) parser.add_argument("--weight_decay", type=float, default=0.05) parser.add_argument("--max_epochs", type=int, default=100) parser.add_argument("--warmup_epochs", type=int, default=5) parser.add_argument("--patience", type=int, default=12) parser.add_argument("--grad_clip", type=float, default=1.0) parser.add_argument("--num_workers", type=int, default=4) parser.add_argument("--device", default="cuda") parser.add_argument("--amp", action=argparse.BooleanOptionalAction, default=True) args = parser.parse_args() if not 0.0 < args.val_fraction < 1.0: parser.error("--val_fraction must be between 0 and 1") if not 0.0 <= args.mask_ratio < 1.0: parser.error("--mask_ratio must be in [0, 1)") return args def select_rows( cache: ExposureCache, train_eids: set[int], val_fraction: float, seed: int ) -> tuple[np.ndarray, np.ndarray]: candidate_eids = np.asarray( sorted(set(map(int, cache.eids)) & train_eids), dtype=np.int64 ) if len(candidate_eids) < 2: raise ValueError("Need at least two training EIDs with cached exposure") rng = np.random.default_rng(seed) rng.shuffle(candidate_eids) n_val = max(1, int(round(len(candidate_eids) * val_fraction))) val_eids = candidate_eids[:n_val] fit_eids = candidate_eids[n_val:] valid_row = np.asarray(cache.row_index, dtype=np.int64) >= 0 fit_event_rows = valid_row & np.isin(cache.eids, fit_eids) val_event_rows = valid_row & np.isin(cache.eids, val_eids) fit_rows = np.unique(np.asarray(cache.row_index[fit_event_rows], dtype=np.int64)) val_rows = np.unique(np.asarray(cache.row_index[val_event_rows], dtype=np.int64)) if len(fit_rows) == 0 or len(val_rows) == 0: raise ValueError("Training/validation exposure rows are empty after filtering") return fit_rows, val_rows def channel_stats( cache: ExposureCache, rows: np.ndarray, chunk_size: int = 256 ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: results = [] for source in (cache.daily, cache.monthly): sums = np.zeros(source.shape[-1], dtype=np.float64) squares = np.zeros_like(sums) counts = np.zeros_like(sums) for start in tqdm(range(0, len(rows), chunk_size), desc="Channel statistics"): values = np.asarray(source[rows[start:start + chunk_size]], dtype=np.float64) finite = np.isfinite(values) clean = np.where(finite, values, 0.0) sums += clean.sum(axis=(0, 1)) squares += np.square(clean).sum(axis=(0, 1)) counts += finite.sum(axis=(0, 1)) mean = sums / np.maximum(counts, 1.0) variance = squares / np.maximum(counts, 1.0) - np.square(mean) std = np.sqrt(np.maximum(variance, 1e-12)) results.extend([mean.astype(np.float32), std.astype(np.float32)]) return tuple(results) def masked_mse( prediction: torch.Tensor, target: torch.Tensor, mask: torch.Tensor ) -> torch.Tensor: error = (prediction - target).square() * mask return error.sum() / mask.sum().clamp_min(1.0) def run_epoch( model: TimesNetExposureAutoencoder, loader: DataLoader, device: torch.device, stats: tuple[torch.Tensor, ...], mask_ratio: float, optimizer: AdamW | None, scaler: torch.amp.GradScaler, grad_clip: float, amp_enabled: bool, ) -> float: training = optimizer is not None model.train(training) total_loss = 0.0 total_samples = 0 daily_mean, daily_std, monthly_mean, monthly_std = stats context = torch.enable_grad if training else torch.no_grad with context(): for batch in tqdm(loader, desc="train" if training else "val", leave=False): daily = batch["daily"].to(device, non_blocking=True) monthly = batch["monthly"].to(device, non_blocking=True) daily_observed = torch.isfinite(daily) monthly_observed = torch.isfinite(monthly) daily = (torch.nan_to_num(daily) - daily_mean) / daily_std monthly = (torch.nan_to_num(monthly) - monthly_mean) / monthly_std daily = daily * daily_observed monthly = monthly * monthly_observed if training and mask_ratio > 0: daily_input_mask = daily_observed & ( torch.rand_like(daily) >= mask_ratio ) monthly_input_mask = monthly_observed & ( torch.rand_like(monthly) >= mask_ratio ) else: daily_input_mask = daily_observed monthly_input_mask = monthly_observed daily_input = daily * daily_input_mask monthly_input = monthly * monthly_input_mask if training: optimizer.zero_grad(set_to_none=True) with torch.autocast( device_type=device.type, dtype=torch.float16, enabled=amp_enabled, ): daily_hat, monthly_hat, _ = model( daily_input, monthly_input, daily_input_mask, monthly_input_mask, ) loss = ( masked_mse(daily_hat, daily, daily_observed) + masked_mse(monthly_hat, monthly, monthly_observed) ) if training: scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip) scaler.step(optimizer) scaler.update() batch_size = daily.size(0) total_loss += float(loss.detach()) * batch_size total_samples += batch_size return total_loss / max(total_samples, 1) def learning_rate(epoch: int, args: argparse.Namespace) -> float: if epoch < args.warmup_epochs: return args.base_lr * (epoch + 1) / max(args.warmup_epochs, 1) progress = (epoch - args.warmup_epochs) / max( args.max_epochs - args.warmup_epochs - 1, 1 ) return args.base_lr * 0.5 * (1.0 + math.cos(math.pi * progress)) def main() -> None: args = parse_args() set_seed(args.seed) device = resolve_device(args.device) configure_torch_for_training(device) run_dir, run_name = create_unique_run_dir( lambda stamp: f"exposure_ae_{stamp}", Path(args.runs_root) ) logger = setup_logging(run_dir) cache = ExposureCache(args.exposure_cache_dir) train_rows, val_rows = select_rows( cache, load_eid_file(args.train_eid_file), args.val_fraction, args.seed ) raw_stats = channel_stats(cache, train_rows) stats = tuple( torch.as_tensor(value, device=device).view(1, 1, -1) for value in raw_stats ) loader_kwargs = dict( batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=device.type == "cuda", persistent_workers=args.num_workers > 0, ) train_loader = DataLoader( ExposureWindowDataset(cache, train_rows), shuffle=True, **loader_kwargs ) val_loader = DataLoader( ExposureWindowDataset(cache, val_rows), shuffle=False, **loader_kwargs ) model = TimesNetExposureAutoencoder( n_embd=args.n_embd, d_model=args.d_model, n_layers=args.n_layers, top_k=args.top_k, n_convnext_blocks=args.n_convnext_blocks, conv_kernel_size=args.conv_kernel_size, mlp_ratio=args.mlp_ratio, dropout=args.dropout, ).to(device) optimizer = AdamW( model.parameters(), lr=args.base_lr, weight_decay=args.weight_decay, betas=(0.9, 0.95), ) amp_enabled = bool(args.amp and device.type == "cuda") scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled) logger.info( f"Run {run_name}: device={device}, train_rows={len(train_rows):,}, " f"val_rows={len(val_rows):,}" ) config = vars(args) | { "train_rows": len(train_rows), "val_rows": len(val_rows), "daily_mean": raw_stats[0].tolist(), "daily_std": raw_stats[1].tolist(), "monthly_mean": raw_stats[2].tolist(), "monthly_std": raw_stats[3].tolist(), } (run_dir / "train_config.json").write_text( json.dumps(config, indent=2), encoding="utf-8" ) best_loss = float("inf") stale_epochs = 0 history = [] for epoch in range(args.max_epochs): lr = learning_rate(epoch, args) for group in optimizer.param_groups: group["lr"] = lr train_loss = run_epoch( model, train_loader, device, stats, args.mask_ratio, optimizer, scaler, args.grad_clip, amp_enabled, ) val_loss = run_epoch( model, val_loader, device, stats, 0.0, None, scaler, args.grad_clip, amp_enabled, ) logger.info( f"Epoch {epoch + 1:03d} | lr={lr:.3e} | " f"train={train_loss:.6f} | val={val_loss:.6f}" ) history.append( {"epoch": epoch + 1, "lr": lr, "train_loss": train_loss, "val_loss": val_loss} ) if val_loss < best_loss: best_loss = val_loss stale_epochs = 0 torch.save( { "model_state_dict": model.state_dict(), "encoder_state_dict": model.encoder.state_dict(), "model_config": { key: config[key] for key in ( "n_embd", "d_model", "n_layers", "top_k", "n_convnext_blocks", "conv_kernel_size", "mlp_ratio", "dropout", ) }, "normalization": { "daily_mean": raw_stats[0], "daily_std": raw_stats[1], "monthly_mean": raw_stats[2], "monthly_std": raw_stats[3], }, "epoch": epoch + 1, "val_loss": val_loss, }, run_dir / "best.pt", ) else: stale_epochs += 1 (run_dir / "history.json").write_text( json.dumps(history, indent=2), encoding="utf-8" ) if stale_epochs >= args.patience: logger.info(f"Early stopping after {epoch + 1} epochs") break logger.info(f"Best validation loss: {best_loss:.6f}") logger.info(f"Checkpoint: {run_dir / 'best.pt'}") if __name__ == "__main__": main()