JiaruiLi/DeepHealthExpo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Optimize exposure autoencoder distributed training

Browse Source
This commit is contained in:
Jiarui Li
2026-07-09 15:44:21 +08:00
parent f7fb6b7718
commit 5a3122c965
3 changed files with 258 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

292
train_exposure_autoencoder.py
View File

@@ -2,14 +2,19 @@
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
from torch.utils.data import DataLoader, Dataset, DistributedSampler
from tqdm import tqdm
from backbones import TimesNetExposureAutoencoder
@@ -51,6 +56,19 @@ def parse_args() -> argparse.Namespace:
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 "
"<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("--seed", type=int, default=42)
parser.add_argument("--n_embd", type=int, default=120)
@@ -82,9 +100,23 @@ def parse_args() -> argparse.Namespace:
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 = [
@@ -145,7 +177,50 @@ def maybe_wrap_data_parallel(
def unwrap_model(model) -> TimesNetExposureAutoencoder:
return model.module if isinstance(model, torch.nn.DataParallel) else model
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(
@@ -170,6 +245,61 @@ def channel_stats(
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:
@@ -187,15 +317,20 @@ def run_epoch(
scaler: torch.amp.GradScaler,
grad_clip: float,
amp_enabled: bool,
show_progress: bool,
) -> float:
training = optimizer is not None
model.train(training)
total_loss = 0.0
total_samples = 0
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):
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)
@@ -239,9 +374,11 @@ def run_epoch(
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)
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:
@@ -255,13 +392,11 @@ def learning_rate(epoch: int, args: argparse.Namespace) -> float:
def main() -> None:
args = parse_args()
set_seed(args.seed)
device = resolve_device(args.device)
device, rank, local_rank, world_size = init_distributed(args)
set_seed(args.seed + rank)
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)
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)
@@ -269,22 +404,65 @@ def main() -> None:
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_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=args.batch_size,
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(
ExposureWindowDataset(cache, train_rows), shuffle=True, **loader_kwargs
train_dataset, sampler=train_sampler,
shuffle=train_sampler is None, **loader_kwargs
)
val_loader = DataLoader(
ExposureWindowDataset(cache, val_rows), shuffle=False, **loader_kwargs
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,
@@ -293,7 +471,16 @@ def main() -> None:
backbone_expansion=args.backbone_expansion,
dropout=args.dropout,
).to(device)
model = maybe_wrap_data_parallel(model, args, device, logger)
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),
@@ -312,24 +499,27 @@ def main() -> None:
"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"
)
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,
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,
scaler, args.grad_clip, amp_enabled, rank == 0,
)
logger.info(
f"Epoch {epoch + 1:03d} | lr={lr:.3e} | "
@@ -342,39 +532,43 @@ def main() -> None:
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",
)
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,
},
"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",
)
run_dir / "best.pt",
)
else:
stale_epochs += 1
(run_dir / "history.json").write_text(
json.dumps(history, indent=2), encoding="utf-8"
)
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__":