382 lines
14 KiB
Python
382 lines
14 KiB
Python
"""Pretrain a lightweight TimesNet 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 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("--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.",
|
|
)
|
|
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:<id>")
|
|
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_backbone_blocks=args.n_backbone_blocks,
|
|
backbone_kernel_size=args.backbone_kernel_size,
|
|
backbone_expansion=args.backbone_expansion,
|
|
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_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
|
|
(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()
|