772 lines
28 KiB
Python
772 lines
28 KiB
Python
"""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", default=None)
|
|
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 "
|
|
"<exposure_cache_dir>/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(
|
|
"--resume_checkpoint",
|
|
default=None,
|
|
help=(
|
|
"Resume training from a run directory or checkpoint. A directory "
|
|
"uses last.pt when present, otherwise best.pt."
|
|
),
|
|
)
|
|
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.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def validate_args(args: argparse.Namespace) -> None:
|
|
if not args.exposure_cache_dir:
|
|
raise ValueError(
|
|
"--exposure_cache_dir is required unless loaded from resume config"
|
|
)
|
|
if not 0.0 <= args.mask_ratio < 1.0:
|
|
raise ValueError("--mask_ratio must be in [0, 1)")
|
|
if args.num_workers > 0 and args.prefetch_factor <= 0:
|
|
raise ValueError("--prefetch_factor must be positive")
|
|
if isinstance(args.gpu_ids, str) and args.gpu_ids:
|
|
try:
|
|
args.gpu_ids = [
|
|
int(part.strip())
|
|
for part in args.gpu_ids.split(",")
|
|
if part.strip()
|
|
]
|
|
except ValueError as exc:
|
|
raise ValueError(
|
|
"--gpu_ids must be a comma-separated list of integers"
|
|
) from exc
|
|
if not args.gpu_ids:
|
|
raise ValueError("--gpu_ids did not contain any valid CUDA device ids")
|
|
args.data_parallel = True
|
|
|
|
|
|
def resolve_resume_checkpoint(resume_path: str | None) -> Path | None:
|
|
if not resume_path:
|
|
return None
|
|
path = Path(resume_path)
|
|
if path.is_dir():
|
|
last_path = path / "last.pt"
|
|
if last_path.is_file():
|
|
return last_path
|
|
best_path = path / "best.pt"
|
|
if best_path.is_file():
|
|
return best_path
|
|
raise FileNotFoundError(
|
|
f"Resume run directory has neither last.pt nor best.pt: {path}"
|
|
)
|
|
if not path.is_file():
|
|
raise FileNotFoundError(f"--resume_checkpoint does not exist: {path}")
|
|
return path
|
|
|
|
|
|
def apply_resume_config(args: argparse.Namespace, resume_checkpoint: Path) -> None:
|
|
config_path = resume_checkpoint.parent / "train_config.json"
|
|
if not config_path.is_file():
|
|
raise FileNotFoundError(
|
|
f"Resume requires train_config.json next to checkpoint: {config_path}"
|
|
)
|
|
config = json.loads(config_path.read_text(encoding="utf-8"))
|
|
resume_value = str(resume_checkpoint)
|
|
for key, value in config.items():
|
|
if hasattr(args, key):
|
|
setattr(args, key, value)
|
|
args.resume_checkpoint = resume_value
|
|
|
|
|
|
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:
|
|
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:
|
|
if args.resume_checkpoint:
|
|
resume_path = Path(args.resume_checkpoint)
|
|
run_dir = resume_path.parent
|
|
run_name = run_dir.name
|
|
else:
|
|
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 autoencoder_checkpoint_payload(
|
|
model,
|
|
optimizer: AdamW,
|
|
scaler: torch.amp.GradScaler,
|
|
config: dict,
|
|
raw_stats: tuple[np.ndarray, ...],
|
|
epoch: int,
|
|
val_loss: float,
|
|
best_loss: float,
|
|
stale_epochs: int,
|
|
history: list[dict],
|
|
include_training_state: bool,
|
|
) -> dict:
|
|
checkpoint_model = unwrap_model(model)
|
|
payload = {
|
|
"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,
|
|
"val_loss": val_loss,
|
|
"best_loss": best_loss,
|
|
"stale_epochs": stale_epochs,
|
|
"history": history,
|
|
}
|
|
if include_training_state:
|
|
payload["optimizer_state_dict"] = optimizer.state_dict()
|
|
payload["scaler_state_dict"] = scaler.state_dict()
|
|
return payload
|
|
|
|
|
|
def save_autoencoder_checkpoint(payload: dict, checkpoint_path: Path) -> None:
|
|
tmp_path = checkpoint_path.with_suffix(checkpoint_path.suffix + ".tmp")
|
|
torch.save(payload, tmp_path)
|
|
tmp_path.replace(checkpoint_path)
|
|
|
|
|
|
def load_resume_checkpoint(
|
|
checkpoint_path: Path,
|
|
model,
|
|
optimizer: AdamW,
|
|
scaler: torch.amp.GradScaler,
|
|
val_loader: DataLoader,
|
|
stats: tuple[torch.Tensor, ...],
|
|
device: torch.device,
|
|
grad_clip: float,
|
|
amp_enabled: bool,
|
|
show_progress: bool,
|
|
logger,
|
|
) -> tuple[int, float, int, list[dict]]:
|
|
checkpoint = torch.load(checkpoint_path, map_location=device)
|
|
if "model_state_dict" not in checkpoint:
|
|
raise KeyError(
|
|
f"Checkpoint does not contain model_state_dict: {checkpoint_path}"
|
|
)
|
|
unwrap_model(model).load_state_dict(checkpoint["model_state_dict"])
|
|
has_training_state = "optimizer_state_dict" in checkpoint
|
|
if has_training_state:
|
|
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
|
|
else:
|
|
logger.warning(
|
|
"Resume checkpoint has no optimizer state; continuing with a fresh "
|
|
"optimizer"
|
|
)
|
|
if "scaler_state_dict" in checkpoint:
|
|
scaler.load_state_dict(checkpoint["scaler_state_dict"])
|
|
elif scaler.is_enabled():
|
|
logger.warning(
|
|
"Resume checkpoint has no AMP scaler state; continuing with a fresh "
|
|
"scaler"
|
|
)
|
|
start_epoch = int(checkpoint.get("epoch", 0))
|
|
history = list(checkpoint.get("history", []))
|
|
history_path = checkpoint_path.parent / "history.json"
|
|
if not history and history_path.is_file():
|
|
history = json.loads(history_path.read_text(encoding="utf-8"))
|
|
|
|
if has_training_state:
|
|
best_loss = float(
|
|
checkpoint.get("best_loss", checkpoint.get("val_loss", float("inf")))
|
|
)
|
|
stale_epochs = int(checkpoint.get("stale_epochs", 0))
|
|
else:
|
|
current_val_loss = run_epoch(
|
|
model, val_loader, device, stats, 0.0, None,
|
|
scaler, grad_clip, amp_enabled, show_progress,
|
|
)
|
|
history_best_epoch = start_epoch
|
|
historical_best = float(checkpoint.get("val_loss", float("inf")))
|
|
for entry in history:
|
|
if "val_loss" not in entry:
|
|
continue
|
|
entry_loss = float(entry["val_loss"])
|
|
if entry_loss < historical_best:
|
|
historical_best = entry_loss
|
|
history_best_epoch = int(entry.get("epoch", len(history)))
|
|
if math.isfinite(historical_best):
|
|
tolerance = max(1e-4, abs(historical_best) * 1e-3)
|
|
if abs(current_val_loss - historical_best) > tolerance:
|
|
logger.warning(
|
|
"Legacy best.pt validation loss differs from history: "
|
|
f"recomputed={current_val_loss:.6f}, "
|
|
f"history_best={historical_best:.6f}"
|
|
)
|
|
best_loss = min(historical_best, current_val_loss)
|
|
if history:
|
|
start_epoch = max(start_epoch, int(history[-1].get("epoch", len(history))))
|
|
if current_val_loss < historical_best:
|
|
stale_epochs = 0
|
|
else:
|
|
stale_epochs = max(0, start_epoch - history_best_epoch)
|
|
logger.info(
|
|
f"Validated legacy checkpoint {checkpoint_path.name}: "
|
|
f"val={current_val_loss:.6f}, historical_best={historical_best:.6f}"
|
|
)
|
|
logger.info(
|
|
f"Resumed from {checkpoint_path} at epoch {start_epoch}; "
|
|
f"best_val={best_loss:.6f}, stale_epochs={stale_epochs}"
|
|
)
|
|
return start_epoch, best_loss, stale_epochs, history
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
resume_checkpoint = resolve_resume_checkpoint(args.resume_checkpoint)
|
|
if resume_checkpoint is not None:
|
|
apply_resume_config(args, resume_checkpoint)
|
|
validate_args(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 and not args.resume_checkpoint:
|
|
(run_dir / "train_config.json").write_text(
|
|
json.dumps(config, indent=2), encoding="utf-8"
|
|
)
|
|
|
|
best_loss = float("inf")
|
|
stale_epochs = 0
|
|
history = []
|
|
start_epoch = 0
|
|
if args.resume_checkpoint:
|
|
start_epoch, best_loss, stale_epochs, history = load_resume_checkpoint(
|
|
Path(args.resume_checkpoint), model, optimizer, scaler,
|
|
val_loader, stats, device, args.grad_clip, amp_enabled,
|
|
rank == 0, logger,
|
|
)
|
|
if start_epoch >= args.max_epochs:
|
|
logger.info(
|
|
f"Resume checkpoint is already at epoch {start_epoch}; "
|
|
f"--max_epochs={args.max_epochs} leaves no remaining epochs"
|
|
)
|
|
for epoch in range(start_epoch, 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:
|
|
save_autoencoder_checkpoint(
|
|
autoencoder_checkpoint_payload(
|
|
model, optimizer, scaler, config, raw_stats,
|
|
epoch + 1, val_loss, best_loss, stale_epochs,
|
|
history, include_training_state=False,
|
|
),
|
|
run_dir / "best.pt",
|
|
)
|
|
else:
|
|
stale_epochs += 1
|
|
if rank == 0:
|
|
save_autoencoder_checkpoint(
|
|
autoencoder_checkpoint_payload(
|
|
model, optimizer, scaler, config, raw_stats,
|
|
epoch + 1, val_loss, best_loss, stale_epochs,
|
|
history, include_training_state=True,
|
|
),
|
|
run_dir / "last.pt",
|
|
)
|
|
(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()
|