"""Pretrain a lightweight TimesNet autoencoder on training-set exposure.""" from __future__ import annotations import argparse import hashlib import json import logging import math import os from pathlib import Path import numpy as np import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel from torch.optim import AdamW from torch.utils.data import DataLoader, Dataset, DistributedSampler 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 lightweight TimesNet 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( "--channel_stats_file", default=None, help=( "Cached channel statistics .npz file. Defaults to " "/train_channel_stats.npz." ), ) parser.add_argument( "--recompute_channel_stats", action="store_true", help="Ignore a compatible statistics cache and recompute it.", ) 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=64) parser.add_argument("--n_layers", type=int, default=2) parser.add_argument("--top_k", type=int, default=2) parser.add_argument("--n_backbone_blocks", type=int, default=1) parser.add_argument("--backbone_kernel_size", type=int, default=5) parser.add_argument("--backbone_expansion", type=float, default=2.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.", ) parser.add_argument( "--ddp_backend", default=None, choices=["nccl", "gloo"], help="DDP backend. Defaults to nccl on CUDA and gloo otherwise.", ) parser.add_argument( "--prefetch_factor", type=int, default=4, help="DataLoader batches prefetched by each worker.", ) args = parser.parse_args() if not 0.0 <= args.mask_ratio < 1.0: parser.error("--mask_ratio must be in [0, 1)") if args.num_workers > 0 and args.prefetch_factor <= 0: parser.error("--prefetch_factor must be positive") 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: if isinstance(model, (torch.nn.DataParallel, DistributedDataParallel)): return model.module return model def init_distributed( args: argparse.Namespace, ) -> tuple[torch.device, int, int, int]: world_size = int(os.environ.get("WORLD_SIZE", "1")) if world_size == 1: return resolve_device(args.device), 0, 0, 1 if args.data_parallel: raise ValueError("--data_parallel cannot be combined with torchrun/DDP") local_rank = int(os.environ["LOCAL_RANK"]) rank = int(os.environ["RANK"]) if not torch.cuda.is_available(): raise ValueError("Multi-process exposure training requires CUDA") torch.cuda.set_device(local_rank) backend = args.ddp_backend or "nccl" dist.init_process_group(backend=backend, init_method="env://") return torch.device("cuda", local_rank), rank, local_rank, world_size def rank_logger(rank: int, run_dir: Path): if rank == 0: return setup_logging(run_dir) logger = logging.getLogger(f"DeepHealth.rank{rank}") logger.handlers.clear() logger.addHandler(logging.NullHandler()) return logger def distributed_run_dir( args: argparse.Namespace, rank: int, world_size: int ) -> tuple[Path, str]: payload: list[str | None] = [None, None] if rank == 0: run_dir, run_name = create_unique_run_dir( lambda stamp: f"exposure_ae_{stamp}", Path(args.runs_root) ) payload = [str(run_dir), run_name] if world_size > 1: dist.broadcast_object_list(payload, src=0) return Path(str(payload[0])), str(payload[1]) 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 eid_set_hash(eids: set[int]) -> str: digest = hashlib.sha256() for eid in sorted(eids): digest.update(f"{eid}\n".encode("ascii")) return digest.hexdigest() def load_or_compute_channel_stats( cache: ExposureCache, rows: np.ndarray, train_eids: set[int], stats_path: Path, recompute: bool, logger, ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: eid_hash = eid_set_hash(train_eids) if stats_path.is_file() and not recompute: try: with np.load(stats_path, allow_pickle=False) as saved: compatible = ( str(saved["train_eid_sha256"].item()) == eid_hash and int(saved["cache_event_rows"].item()) == len(cache.eids) and int(saved["train_window_rows"].item()) == len(rows) ) if compatible: logger.info(f"Loading channel statistics from {stats_path}") return ( saved["daily_mean"].astype(np.float32), saved["daily_std"].astype(np.float32), saved["monthly_mean"].astype(np.float32), saved["monthly_std"].astype(np.float32), ) logger.info("Channel statistics cache is stale; recomputing") except (KeyError, OSError, ValueError) as exc: logger.warning( f"Could not read channel statistics cache ({exc}); recomputing" ) logger.info("Computing channel statistics from training exposure") stats = channel_stats(cache, rows) stats_path.parent.mkdir(parents=True, exist_ok=True) np.savez( stats_path, daily_mean=stats[0], daily_std=stats[1], monthly_mean=stats[2], monthly_std=stats[3], train_eid_sha256=np.asarray(eid_hash), cache_event_rows=np.asarray(len(cache.eids), dtype=np.int64), train_window_rows=np.asarray(len(rows), dtype=np.int64), ) logger.info(f"Saved channel statistics to {stats_path}") return stats 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, show_progress: bool, ) -> float: training = optimizer is not None model.train(training) loss_accumulator = torch.zeros(2, device=device, dtype=torch.float64) 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, disable=not show_progress, ): 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) loss_accumulator[0] += loss.detach().double() * batch_size loss_accumulator[1] += batch_size if dist.is_initialized(): dist.all_reduce(loss_accumulator, op=dist.ReduceOp.SUM) return float((loss_accumulator[0] / loss_accumulator[1].clamp_min(1)).item()) 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() device, rank, local_rank, world_size = init_distributed(args) set_seed(args.seed + rank) configure_torch_for_training(device) run_dir, run_name = distributed_run_dir(args, rank, world_size) logger = rank_logger(rank, 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") stats_path = ( Path(args.channel_stats_file) if args.channel_stats_file else Path(args.exposure_cache_dir) / "train_channel_stats.npz" ) if rank == 0: raw_stats = load_or_compute_channel_stats( cache, train_rows, train_eids, stats_path, args.recompute_channel_stats, logger, ) if world_size > 1: dist.barrier() if rank != 0: raw_stats = load_or_compute_channel_stats( cache, train_rows, train_eids, stats_path, False, logger ) stats = tuple( torch.as_tensor(value, device=device).view(1, 1, -1) for value in raw_stats ) if args.batch_size % world_size != 0: raise ValueError( f"--batch_size={args.batch_size} must be divisible by " f"DDP world size {world_size}" ) local_batch_size = args.batch_size // world_size loader_kwargs = dict( batch_size=local_batch_size, num_workers=args.num_workers, pin_memory=device.type == "cuda", persistent_workers=args.num_workers > 0, ) if args.num_workers > 0: loader_kwargs["prefetch_factor"] = args.prefetch_factor train_dataset = ExposureWindowDataset(cache, train_rows) val_dataset = ExposureWindowDataset(cache, val_rows) train_sampler = ( DistributedSampler( train_dataset, num_replicas=world_size, rank=rank, shuffle=True, seed=args.seed, ) if world_size > 1 else None ) val_sampler = ( DistributedSampler( val_dataset, num_replicas=world_size, rank=rank, shuffle=False ) if world_size > 1 else None ) train_loader = DataLoader( train_dataset, sampler=train_sampler, shuffle=train_sampler is None, **loader_kwargs ) val_loader = DataLoader( val_dataset, sampler=val_sampler, 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_backbone_blocks=args.n_backbone_blocks, backbone_kernel_size=args.backbone_kernel_size, backbone_expansion=args.backbone_expansion, dropout=args.dropout, ).to(device) if world_size > 1: model = DistributedDataParallel( model, device_ids=[local_rank], output_device=local_rank ) logger.info( f"Using DDP with {world_size} processes; " f"global_batch={args.batch_size}, per_gpu_batch={local_batch_size}" ) else: 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(), } if rank == 0: (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): if train_sampler is not None: train_sampler.set_epoch(epoch) 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, rank == 0, ) val_loss = run_epoch( model, val_loader, device, stats, 0.0, None, scaler, args.grad_clip, amp_enabled, rank == 0, ) 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 if rank == 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_backbone_blocks", "backbone_kernel_size", "backbone_expansion", "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 if rank == 0: (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 dist.is_initialized(): dist.destroy_process_group() if __name__ == "__main__": main()