672 lines
24 KiB
Python
672 lines
24 KiB
Python
"""
|
|
Train DeepHealth with query-conditioned all-future supervision.
|
|
|
|
Training samples are patient-level. For each patient and each __getitem__ call,
|
|
AllFutureHealthDataset randomly samples a query time t_query, uses events at or
|
|
before t_query as history, and uses events after t_query as the future target set.
|
|
|
|
All splits use the same patient/interval/time-uniform query distribution.
|
|
Validation/test keep one deterministic query draw per patient.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import math
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, Dict
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torch.nn.utils import clip_grad_norm_
|
|
from torch.optim import AdamW
|
|
from torch.utils.data import DataLoader, RandomSampler
|
|
from tqdm.auto import tqdm
|
|
|
|
from dataset import (
|
|
DISEASE_HISTORY_MODES,
|
|
DISEASE_HISTORY_MODE_TIMED,
|
|
AllFutureHealthDataset,
|
|
all_future_collate_fn,
|
|
)
|
|
from losses import build_loss
|
|
from model_architectures import (
|
|
DEFAULT_MODEL_ARCHITECTURE,
|
|
SUPPORTED_MODEL_ARCHITECTURES,
|
|
)
|
|
from models import DeepHealth
|
|
from targets import NO_EVENT_IDX, PAD_IDX, RESERVED_IDX
|
|
from train_util import (
|
|
AllFutureBaselineStats,
|
|
ContinuousRobustScalerStats,
|
|
configure_torch_for_training,
|
|
create_unique_run_dir,
|
|
format_extra_info_types,
|
|
fit_continuous_robust_scaler,
|
|
fit_all_future_baseline,
|
|
get_lr,
|
|
get_model_parameter_counts,
|
|
load_extra_info_types_file,
|
|
move_batch_to_device,
|
|
resolve_device,
|
|
save_checkpoint,
|
|
save_config,
|
|
set_optimizer_lr,
|
|
set_seed,
|
|
setup_logging,
|
|
split_all_future_datasets,
|
|
split_all_future_datasets_by_eid_files,
|
|
)
|
|
|
|
|
|
MODEL_INPUT_KEYS = (
|
|
"event_seq",
|
|
"time_seq",
|
|
"sex",
|
|
"padding_mask",
|
|
"t_query",
|
|
"other_type",
|
|
"other_value",
|
|
"other_value_kind",
|
|
"other_time",
|
|
)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Train DeepHealth with all-future supervision")
|
|
|
|
parser.add_argument("--data_prefix", type=str, default="ukb")
|
|
parser.add_argument("--labels_file", type=str, default="labels.csv")
|
|
parser.add_argument("--runs_root", type=str, default="runs")
|
|
parser.add_argument("--seed", type=int, default=42)
|
|
parser.add_argument("--extra_info_types_file", type=str, default=None)
|
|
parser.add_argument(
|
|
"--disease_history_mode",
|
|
type=str,
|
|
default=DISEASE_HISTORY_MODE_TIMED,
|
|
choices=DISEASE_HISTORY_MODES,
|
|
help=(
|
|
"timed=real disease times; ordered=chronological disease order with "
|
|
"ordinal positions; set=unordered disease set with no disease time"
|
|
),
|
|
)
|
|
|
|
parser.add_argument("--train_ratio", type=float, default=0.7)
|
|
parser.add_argument("--val_ratio", type=float, default=0.15)
|
|
parser.add_argument("--test_ratio", type=float, default=0.15)
|
|
parser.add_argument("--train_eid_file", type=str, default="ukb_train_eid.csv")
|
|
parser.add_argument("--val_eid_file", type=str, default="ukb_val_eid.csv")
|
|
parser.add_argument("--test_eid_file", type=str, default="ukb_test_eid.csv")
|
|
parser.add_argument("--min_history_events", type=int, default=1)
|
|
parser.add_argument("--min_future_events", type=int, default=1)
|
|
parser.add_argument("--validation_query_seed", type=int, default=None)
|
|
|
|
parser.add_argument("--n_embd", type=int, default=120)
|
|
parser.add_argument("--n_head", type=int, default=10)
|
|
parser.add_argument("--n_layer", type=int, default=12)
|
|
parser.add_argument("--n_bins", type=int, default=16)
|
|
parser.add_argument("--extra_pool_reduce", type=str, default="mean",
|
|
choices=["mean", "sum"])
|
|
parser.add_argument("--time_mode", type=str, default="relative",
|
|
choices=["relative", "absolute"])
|
|
parser.add_argument("--dist_mode", type=str, default="exponential",
|
|
choices=["exponential", "weibull"])
|
|
parser.add_argument("--dropout", type=float, default=0.0)
|
|
parser.add_argument(
|
|
"--risk_head_bias",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=True,
|
|
help=(
|
|
"Initialize a learnable all-future output bias from training-only "
|
|
"marginal first-onset rates"
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--model_architecture",
|
|
type=str,
|
|
default=DEFAULT_MODEL_ARCHITECTURE,
|
|
choices=SUPPORTED_MODEL_ARCHITECTURES,
|
|
)
|
|
|
|
parser.add_argument("--batch_size", type=int, default=128)
|
|
parser.add_argument("--base_lr", type=float, default=3e-4)
|
|
parser.add_argument("--weight_decay", type=float, default=0.1)
|
|
parser.add_argument("--betas", type=float, nargs=2, default=(0.9, 0.99))
|
|
parser.add_argument("--grad_clip", type=float, default=1.0)
|
|
parser.add_argument("--max_epochs", type=int, default=200)
|
|
parser.add_argument("--warmup_epochs", type=int, default=10)
|
|
parser.add_argument("--patience", type=int, default=15)
|
|
parser.add_argument("--min_lr_ratio", type=float, default=0.1)
|
|
parser.add_argument("--num_workers", type=int, default=4)
|
|
parser.add_argument("--device", type=str, default="cuda")
|
|
parser.add_argument("--progress_interval", type=int, default=20)
|
|
|
|
args = parser.parse_args()
|
|
if args.min_history_events < 1:
|
|
raise ValueError("min_history_events must be >= 1")
|
|
if args.min_future_events < 1:
|
|
raise ValueError("min_future_events must be >= 1")
|
|
use_eid_split = all(
|
|
getattr(args, name)
|
|
for name in ("train_eid_file", "val_eid_file", "test_eid_file")
|
|
)
|
|
if not use_eid_split and not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0):
|
|
raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0")
|
|
if args.validation_query_seed is None:
|
|
args.validation_query_seed = int(args.seed)
|
|
args.extra_info_types = (
|
|
load_extra_info_types_file(args.extra_info_types_file)
|
|
if args.extra_info_types_file is not None
|
|
else None
|
|
)
|
|
if args.disease_history_mode != DISEASE_HISTORY_MODE_TIMED:
|
|
expected = {
|
|
"model_architecture": "traj_mixer_v5",
|
|
"time_mode": "relative",
|
|
"dist_mode": "weibull",
|
|
}
|
|
mismatches = [
|
|
f"{name}={getattr(args, name)!r} (expected {value!r})"
|
|
for name, value in expected.items()
|
|
if getattr(args, name) != value
|
|
]
|
|
if args.extra_info_types != []:
|
|
mismatches.append(
|
|
"extra_info_types must be [] via extra_info_types_none.txt"
|
|
)
|
|
if mismatches:
|
|
raise ValueError(
|
|
f"disease_history_mode={args.disease_history_mode!r} is reserved "
|
|
"for the no-extra TrajMixer + all_future + relative + Weibull "
|
|
"ablation; " + "; ".join(mismatches)
|
|
)
|
|
return args
|
|
|
|
|
|
def build_model(
|
|
args: argparse.Namespace,
|
|
dataset: AllFutureHealthDataset,
|
|
scaler_stats: ContinuousRobustScalerStats,
|
|
baseline_stats: AllFutureBaselineStats | None,
|
|
) -> DeepHealth:
|
|
if tuple(int(x) for x in dataset.cont_type_ids) != scaler_stats.cont_type_ids:
|
|
raise ValueError(
|
|
"RobustScale statistics are not aligned with dataset.cont_type_ids"
|
|
)
|
|
center = scaler_stats.center if dataset.n_cont_types > 0 else None
|
|
scale = scaler_stats.scale if dataset.n_cont_types > 0 else None
|
|
if args.risk_head_bias and baseline_stats is None:
|
|
raise ValueError("baseline_stats is required when risk_head_bias is enabled")
|
|
return DeepHealth(
|
|
vocab_size=dataset.vocab_size,
|
|
n_embd=args.n_embd,
|
|
n_head=args.n_head,
|
|
n_layer=args.n_layer,
|
|
n_types=dataset.n_types,
|
|
n_cont_types=dataset.n_cont_types,
|
|
n_categories=dataset.n_categories,
|
|
cont_type_ids=dataset.cont_type_ids,
|
|
n_bins=args.n_bins,
|
|
continuous_value_center=center,
|
|
continuous_value_scale=scale,
|
|
extra_pool_reduce=args.extra_pool_reduce,
|
|
target_mode="all_future",
|
|
time_mode=args.time_mode,
|
|
dist_mode=args.dist_mode,
|
|
dropout=args.dropout,
|
|
model_architecture=args.model_architecture,
|
|
risk_head_bias=args.risk_head_bias,
|
|
risk_head_bias_init=(
|
|
baseline_stats.bias
|
|
if args.risk_head_bias and baseline_stats is not None
|
|
else None
|
|
),
|
|
)
|
|
|
|
|
|
def build_criterion(args: argparse.Namespace):
|
|
ignored_idx = {PAD_IDX, RESERVED_IDX, NO_EVENT_IDX}
|
|
if args.dist_mode == "exponential":
|
|
return build_loss("exponential", ignored_idx=ignored_idx)
|
|
if args.dist_mode == "weibull":
|
|
return build_loss("weibull", ignored_idx=ignored_idx)
|
|
raise ValueError(f"Unknown dist_mode: {args.dist_mode}")
|
|
|
|
|
|
def compute_all_future_loss(
|
|
args: argparse.Namespace,
|
|
model: DeepHealth,
|
|
criterion,
|
|
batch: Dict[str, torch.Tensor],
|
|
device: torch.device,
|
|
) -> torch.Tensor:
|
|
required_keys = set(MODEL_INPUT_KEYS)
|
|
required_keys.update(("future_targets", "future_dt", "exposure"))
|
|
batch = move_batch_to_device(
|
|
{key: batch[key] for key in required_keys},
|
|
device,
|
|
)
|
|
|
|
hidden = model(
|
|
event_seq=batch["event_seq"],
|
|
time_seq=batch["time_seq"],
|
|
sex=batch["sex"],
|
|
padding_mask=batch["padding_mask"],
|
|
t_query=batch["t_query"],
|
|
other_type=batch["other_type"],
|
|
other_value=batch["other_value"],
|
|
other_value_kind=batch["other_value_kind"],
|
|
other_time=batch["other_time"],
|
|
)
|
|
logits = model.calc_risk(hidden)
|
|
|
|
if args.dist_mode == "exponential":
|
|
loss = criterion(
|
|
logits=logits,
|
|
targets=batch["future_targets"],
|
|
exposure=batch["exposure"],
|
|
dt=batch["future_dt"],
|
|
history=batch["event_seq"],
|
|
)
|
|
elif args.dist_mode == "weibull":
|
|
loss = criterion(
|
|
logits=logits,
|
|
weibull_rho=model.calc_weibull_rho(hidden),
|
|
targets=batch["future_targets"],
|
|
dt=batch["future_dt"],
|
|
exposure=batch["exposure"],
|
|
history=batch["event_seq"],
|
|
)
|
|
else:
|
|
raise ValueError(f"Unknown dist_mode: {args.dist_mode}")
|
|
|
|
if not torch.isfinite(loss):
|
|
raise RuntimeError(f"Loss is not finite: {float(loss.detach().cpu())}")
|
|
return loss
|
|
|
|
|
|
def run_epoch(
|
|
logger: logging.Logger,
|
|
args: argparse.Namespace,
|
|
model: DeepHealth,
|
|
criterion,
|
|
loader: DataLoader,
|
|
optimizer: AdamW | None,
|
|
device: torch.device,
|
|
is_train: bool,
|
|
) -> float:
|
|
model.train(is_train)
|
|
total = torch.zeros((), device=device)
|
|
n_batches = 0
|
|
skipped = 0
|
|
desc = "train" if is_train else "val"
|
|
progress_interval = max(1, int(args.progress_interval))
|
|
|
|
progress = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True)
|
|
for batch_idx, batch in enumerate(progress):
|
|
try:
|
|
loss = compute_all_future_loss(args, model, criterion, batch, device)
|
|
if is_train:
|
|
if optimizer is None:
|
|
raise ValueError("optimizer is required for training")
|
|
optimizer.zero_grad(set_to_none=True)
|
|
loss.backward()
|
|
if args.grad_clip > 0:
|
|
clip_grad_norm_(model.parameters(), args.grad_clip)
|
|
optimizer.step()
|
|
|
|
total = total + loss.detach()
|
|
n_batches += 1
|
|
if (batch_idx + 1) % progress_interval == 0:
|
|
avg = total / max(1, n_batches)
|
|
progress.set_postfix(
|
|
loss=f"{float(loss.detach().cpu()):.4f}",
|
|
avg=f"{float(avg.detach().cpu()):.4f}",
|
|
skipped=skipped,
|
|
)
|
|
except RuntimeError as exc:
|
|
if "Loss is not finite" not in str(exc):
|
|
raise
|
|
skipped += 1
|
|
logger.warning(f"Batch {batch_idx} skipped: {str(exc)[:120]}")
|
|
|
|
if skipped:
|
|
logger.info(f"Skipped {skipped} batches due to non-finite loss")
|
|
return float((total / max(1, n_batches)).detach().cpu()) if n_batches else float("inf")
|
|
|
|
|
|
def build_metadata(
|
|
args: argparse.Namespace,
|
|
dataset: AllFutureHealthDataset,
|
|
run_name: str,
|
|
train_subset,
|
|
val_subset,
|
|
test_subset,
|
|
scaler_stats: ContinuousRobustScalerStats,
|
|
baseline_stats: AllFutureBaselineStats | None,
|
|
) -> Dict[str, Any]:
|
|
scaler_metadata = scaler_stats.as_metadata()
|
|
return {
|
|
"run_name": run_name,
|
|
"dataset_class": "AllFutureHealthDataset",
|
|
"collate_fn": "all_future_collate_fn",
|
|
"model_class": "DeepHealth",
|
|
"model_architecture": args.model_architecture,
|
|
"model_target_mode": "all_future",
|
|
"target_mode": "all_future",
|
|
"event_stream_version": "disease_death_only_v1",
|
|
"uses_assessment_event_token": False,
|
|
"dist_mode": args.dist_mode,
|
|
"disease_history_mode": args.disease_history_mode,
|
|
"all_future_min_history_events": int(args.min_history_events),
|
|
"all_future_min_future_events": int(args.min_future_events),
|
|
"all_future_validation_query_seed": int(args.validation_query_seed),
|
|
"all_future_query_distribution": "patient_interval_time_uniform",
|
|
"all_future_queries_per_validation_patient": 1,
|
|
"all_future_likelihood": "first_onset_survival_v2",
|
|
"all_future_prevalent_outcomes_excluded": True,
|
|
"risk_head_bias": bool(args.risk_head_bias),
|
|
"risk_head_bias_weight_decay": 0.0 if args.risk_head_bias else None,
|
|
"risk_head_baseline": (
|
|
baseline_stats.as_metadata()
|
|
if args.risk_head_bias and baseline_stats is not None
|
|
else {"method": "none"}
|
|
),
|
|
"extra_info_types_file": (
|
|
Path(args.extra_info_types_file).name
|
|
if args.extra_info_types_file is not None
|
|
else None
|
|
),
|
|
"extra_info_types": [int(x) for x in dataset.extra_info_types],
|
|
"continuous_value_scaling": "robust",
|
|
"continuous_value_scaler": scaler_metadata,
|
|
"dataset_metadata": {
|
|
"vocab_size": int(dataset.vocab_size),
|
|
"n_types": int(dataset.n_types),
|
|
"n_cont_types": int(dataset.n_cont_types),
|
|
"n_categories": int(dataset.n_categories),
|
|
"cont_type_ids": [int(x) for x in dataset.cont_type_ids],
|
|
"extra_info_types": [int(x) for x in dataset.extra_info_types],
|
|
"event_stream_version": "disease_death_only_v1",
|
|
"uses_assessment_event_token": False,
|
|
},
|
|
"split_sizes": {
|
|
"train": int(len(train_subset)),
|
|
"val": int(len(val_subset)),
|
|
"test": int(len(test_subset)),
|
|
},
|
|
"resolved_readout_name": "none",
|
|
"resolved_loss_name": args.dist_mode,
|
|
}
|
|
|
|
|
|
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 timestamp: (
|
|
(
|
|
""
|
|
if args.disease_history_mode == DISEASE_HISTORY_MODE_TIMED
|
|
else f"{args.disease_history_mode}_"
|
|
)
|
|
+ f"{args.time_mode}_{args.dist_mode}_"
|
|
f"all_future_pure_disease_{timestamp}"
|
|
),
|
|
runs_root=Path(args.runs_root) / args.model_architecture,
|
|
)
|
|
logger = setup_logging(run_dir)
|
|
|
|
logger.info(f"Starting all-future training run: {run_name}")
|
|
logger.info(f"Device: {device}")
|
|
logger.info(f"Model architecture: {args.model_architecture}")
|
|
logger.info(f"Disease history mode: {args.disease_history_mode}")
|
|
logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}")
|
|
logger.info("Continuous value scaling: RobustScale (required)")
|
|
|
|
logger.info("Loading all-future datasets...")
|
|
train_dataset = AllFutureHealthDataset(
|
|
data_prefix=args.data_prefix,
|
|
labels_file=args.labels_file,
|
|
split="train",
|
|
min_history_events=args.min_history_events,
|
|
min_future_events=args.min_future_events,
|
|
validation_query_seed=args.validation_query_seed,
|
|
extra_info_types=args.extra_info_types,
|
|
disease_history_mode=args.disease_history_mode,
|
|
)
|
|
val_dataset = AllFutureHealthDataset(
|
|
data_prefix=args.data_prefix,
|
|
labels_file=args.labels_file,
|
|
split="valid",
|
|
min_history_events=args.min_history_events,
|
|
min_future_events=args.min_future_events,
|
|
validation_query_seed=args.validation_query_seed,
|
|
extra_info_types=args.extra_info_types,
|
|
disease_history_mode=args.disease_history_mode,
|
|
)
|
|
test_dataset = AllFutureHealthDataset(
|
|
data_prefix=args.data_prefix,
|
|
labels_file=args.labels_file,
|
|
split="test",
|
|
min_history_events=args.min_history_events,
|
|
min_future_events=args.min_future_events,
|
|
validation_query_seed=args.validation_query_seed,
|
|
extra_info_types=args.extra_info_types,
|
|
disease_history_mode=args.disease_history_mode,
|
|
)
|
|
if args.train_eid_file and args.val_eid_file and args.test_eid_file:
|
|
train_subset, val_subset, test_subset = split_all_future_datasets_by_eid_files(
|
|
train_dataset=train_dataset,
|
|
val_dataset=val_dataset,
|
|
test_dataset=test_dataset,
|
|
train_eid_file=args.train_eid_file,
|
|
val_eid_file=args.val_eid_file,
|
|
test_eid_file=args.test_eid_file,
|
|
)
|
|
logger.info(
|
|
"Using eid split files: "
|
|
f"train={args.train_eid_file}, val={args.val_eid_file}, test={args.test_eid_file}"
|
|
)
|
|
else:
|
|
train_subset, val_subset, test_subset = split_all_future_datasets(
|
|
train_dataset=train_dataset,
|
|
val_dataset=val_dataset,
|
|
test_dataset=test_dataset,
|
|
train_ratio=args.train_ratio,
|
|
val_ratio=args.val_ratio,
|
|
test_ratio=args.test_ratio,
|
|
seed=args.seed,
|
|
)
|
|
logger.info(
|
|
f"Using random ratio split: train={args.train_ratio}, "
|
|
f"val={args.val_ratio}, test={args.test_ratio}, seed={args.seed}"
|
|
)
|
|
logger.info(
|
|
f"Patients/queries: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
|
|
)
|
|
|
|
if train_dataset.n_cont_types > 0:
|
|
logger.info(
|
|
"Fitting continuous RobustScaler on the complete training subset: "
|
|
f"patients={len(train_subset):,}, features={train_dataset.n_cont_types}"
|
|
)
|
|
scaler_stats = fit_continuous_robust_scaler(
|
|
train_dataset,
|
|
train_subset,
|
|
)
|
|
if train_dataset.n_cont_types > 0:
|
|
logger.info(
|
|
"Continuous RobustScaler fitted: "
|
|
f"observations={int(scaler_stats.observation_count.sum()):,}, "
|
|
f"min_per_feature={int(scaler_stats.observation_count.min()):,}, "
|
|
f"max_per_feature={int(scaler_stats.observation_count.max()):,}"
|
|
)
|
|
|
|
baseline_stats = None
|
|
if args.risk_head_bias:
|
|
logger.info(
|
|
"Fitting all-future output baseline on one seeded query per "
|
|
f"training patient: patients={len(train_subset):,}"
|
|
)
|
|
baseline_stats = fit_all_future_baseline(
|
|
train_dataset,
|
|
train_subset,
|
|
seed=args.seed,
|
|
)
|
|
baseline_summary = baseline_stats.as_metadata()
|
|
rate_summary = baseline_summary["rate_per_year"]
|
|
logger.info(
|
|
"All-future baseline rates/year: "
|
|
f"min={rate_summary['min']:.8g}, "
|
|
f"median={rate_summary['median']:.8g}, "
|
|
f"max={rate_summary['max']:.8g}, "
|
|
f"first_onsets={baseline_summary['observed_first_onsets']:,}"
|
|
)
|
|
|
|
train_loader = DataLoader(
|
|
train_subset,
|
|
batch_size=args.batch_size,
|
|
sampler=RandomSampler(train_subset, generator=torch.Generator().manual_seed(args.seed)),
|
|
collate_fn=all_future_collate_fn,
|
|
num_workers=args.num_workers,
|
|
pin_memory=device.type == "cuda",
|
|
persistent_workers=args.num_workers > 0,
|
|
prefetch_factor=2 if args.num_workers > 0 else None,
|
|
)
|
|
val_loader = DataLoader(
|
|
val_subset,
|
|
batch_size=args.batch_size,
|
|
shuffle=False,
|
|
collate_fn=all_future_collate_fn,
|
|
num_workers=args.num_workers,
|
|
pin_memory=device.type == "cuda",
|
|
persistent_workers=args.num_workers > 0,
|
|
prefetch_factor=2 if args.num_workers > 0 else None,
|
|
)
|
|
test_loader = DataLoader(
|
|
test_subset,
|
|
batch_size=args.batch_size,
|
|
shuffle=False,
|
|
collate_fn=all_future_collate_fn,
|
|
num_workers=args.num_workers,
|
|
pin_memory=device.type == "cuda",
|
|
persistent_workers=args.num_workers > 0,
|
|
prefetch_factor=2 if args.num_workers > 0 else None,
|
|
)
|
|
|
|
model = build_model(
|
|
args,
|
|
train_dataset,
|
|
scaler_stats=scaler_stats,
|
|
baseline_stats=baseline_stats,
|
|
).to(device)
|
|
parameter_counts = get_model_parameter_counts(model)
|
|
logger.info(
|
|
"Model parameters: "
|
|
f"total={parameter_counts['model_parameter_count']:,}, "
|
|
f"trainable={parameter_counts['trainable_parameter_count']:,}"
|
|
)
|
|
if model.risk_head.bias is None:
|
|
optimizer_parameters = model.parameters()
|
|
else:
|
|
optimizer_parameters = [
|
|
{
|
|
"params": [
|
|
parameter
|
|
for parameter in model.parameters()
|
|
if parameter is not model.risk_head.bias
|
|
],
|
|
"weight_decay": args.weight_decay,
|
|
},
|
|
{
|
|
"params": [model.risk_head.bias],
|
|
"weight_decay": 0.0,
|
|
},
|
|
]
|
|
optimizer = AdamW(
|
|
optimizer_parameters,
|
|
lr=args.base_lr,
|
|
betas=tuple(args.betas),
|
|
weight_decay=args.weight_decay,
|
|
)
|
|
criterion = build_criterion(args)
|
|
adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)
|
|
|
|
train_metadata = build_metadata(
|
|
args,
|
|
train_dataset,
|
|
run_name,
|
|
train_subset,
|
|
val_subset,
|
|
test_subset,
|
|
scaler_stats,
|
|
baseline_stats,
|
|
)
|
|
train_metadata.update(parameter_counts)
|
|
save_config(
|
|
args,
|
|
run_dir / "train_config.json",
|
|
extra=train_metadata,
|
|
)
|
|
|
|
best_val = float("inf")
|
|
patience = 0
|
|
history = []
|
|
best_model_path = run_dir / "best_model.pt"
|
|
start = time.time()
|
|
|
|
for epoch in range(args.max_epochs):
|
|
lr = get_lr(epoch, args, adaptive_lr)
|
|
set_optimizer_lr(optimizer, lr)
|
|
|
|
train_loss = run_epoch(logger, args, model, criterion, train_loader, optimizer, device, True)
|
|
with torch.no_grad():
|
|
val_loss = run_epoch(logger, args, model, criterion, val_loader, None, device, False)
|
|
|
|
is_best = val_loss < best_val
|
|
if is_best:
|
|
best_val = val_loss
|
|
patience = 0
|
|
save_checkpoint(model, best_model_path)
|
|
else:
|
|
patience += 1
|
|
|
|
logger.info(
|
|
f"Epoch {epoch + 1}/{args.max_epochs} | lr={lr:.6f} | "
|
|
f"train_loss={train_loss:.6f} | val_loss={val_loss:.6f} | "
|
|
f"best_val_loss={best_val:.6f} | patience={patience}/{args.patience} | "
|
|
f"elapsed={time.time() - start:.1f}s"
|
|
)
|
|
history.append({
|
|
"epoch": epoch + 1,
|
|
"lr": lr,
|
|
"train_loss": train_loss,
|
|
"val_loss": val_loss,
|
|
"best_val_loss": best_val,
|
|
"is_best": int(is_best),
|
|
})
|
|
if patience >= args.patience:
|
|
logger.info(f"Early stopping triggered at epoch {epoch + 1}")
|
|
break
|
|
|
|
with (run_dir / "history.json").open("w", encoding="utf-8") as f:
|
|
json.dump(history, f, indent=2)
|
|
|
|
logger.info("Evaluating best model on all-future test queries...")
|
|
model.load_state_dict(torch.load(best_model_path, map_location=device))
|
|
with torch.no_grad():
|
|
test_loss = run_epoch(logger, args, model, criterion, test_loader, None, device, False)
|
|
logger.info(f"Test loss: {test_loss:.6f}")
|
|
logger.info(f"Best checkpoint: {best_model_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|