"""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("--val_eid_file", default="ukb_val_eid.csv") parser.add_argument("--runs_root", default="runs") parser.add_argument("--seed", type=int, default=42) 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) parser.add_argument( "--data_parallel", action="store_true", help="Use torch.nn.DataParallel across multiple CUDA devices.", ) parser.add_argument( "--gpu_ids", default=None, help="Comma-separated CUDA device ids for --data_parallel, e.g. 0,1,2,3.", ) args = parser.parse_args() if not 0.0 <= args.mask_ratio < 1.0: parser.error("--mask_ratio must be in [0, 1)") if args.gpu_ids: try: args.gpu_ids = [ int(part.strip()) for part in args.gpu_ids.split(",") if part.strip() ] except ValueError as exc: parser.error("--gpu_ids must be a comma-separated list of integers") if not args.gpu_ids: parser.error("--gpu_ids did not contain any valid CUDA device ids") args.data_parallel = True return args def select_rows(cache: ExposureCache, eids: set[int], split: str) -> np.ndarray: valid_row = np.asarray(cache.row_index, dtype=np.int64) >= 0 selected_events = valid_row & np.isin(cache.eids, np.fromiter(eids, np.int64)) rows = np.unique( np.asarray(cache.row_index[selected_events], dtype=np.int64) ) if len(rows) == 0: raise ValueError(f"{split} exposure rows are empty after EID filtering") return rows def maybe_wrap_data_parallel( model: TimesNetExposureAutoencoder, args: argparse.Namespace, device: torch.device, logger, ): if not args.data_parallel: return model if device.type != "cuda": raise ValueError("--data_parallel requires --device cuda or cuda:") if not torch.cuda.is_available() or torch.cuda.device_count() < 2: raise ValueError("--data_parallel requires at least two CUDA devices") primary = ( int(device.index) if device.index is not None else int(torch.cuda.current_device()) ) device_ids = ( args.gpu_ids if args.gpu_ids else list(range(torch.cuda.device_count())) ) device_ids = [primary, *[idx for idx in device_ids if idx != primary]] if len(device_ids) < 2: raise ValueError("--data_parallel needs at least two device ids") if any(idx < 0 or idx >= torch.cuda.device_count() for idx in device_ids): raise ValueError(f"CUDA device id is out of range: {device_ids}") logger.info(f"Using DataParallel on CUDA devices: {device_ids}") return torch.nn.DataParallel( model, device_ids=device_ids, output_device=primary ) def unwrap_model(model) -> TimesNetExposureAutoencoder: return model.module if isinstance(model, torch.nn.DataParallel) else model 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, 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_eids = load_eid_file(args.train_eid_file) val_eids = load_eid_file(args.val_eid_file) if train_eids & val_eids: raise ValueError("train and validation EID files must be disjoint") train_rows = select_rows(cache, train_eids, "Training") val_rows = select_rows(cache, val_eids, "Validation") 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) model = maybe_wrap_data_parallel(model, args, device, logger) 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 checkpoint_model = unwrap_model(model) torch.save( { "model_state_dict": checkpoint_model.state_dict(), "encoder_state_dict": checkpoint_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()