Initial DeepHealthExpo next-token codebase
This commit is contained in:
667
train_next_step.py
Normal file
667
train_next_step.py
Normal file
@@ -0,0 +1,667 @@
|
||||
"""
|
||||
Train DeepHealth with next-token / next-time-point supervision.
|
||||
|
||||
The next-step dataset uses observed event histories, including CHECKUP state
|
||||
tokens, plus optional gap <NO_EVENT> imputation. UTS training reads out only
|
||||
same-time group ends.
|
||||
"""
|
||||
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 HealthDataset, collate_fn
|
||||
from losses import build_loss
|
||||
from models import DeepHealth, DeepHealthOutput
|
||||
from readouts import build_readout
|
||||
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
|
||||
from train_util import (
|
||||
configure_torch_for_training,
|
||||
create_unique_run_dir,
|
||||
format_extra_info_types,
|
||||
load_extra_info_types_file,
|
||||
resolve_device,
|
||||
save_checkpoint,
|
||||
save_config,
|
||||
set_optimizer_lr,
|
||||
set_seed,
|
||||
setup_logging,
|
||||
split_dataset,
|
||||
split_dataset_by_eid_files,
|
||||
)
|
||||
|
||||
|
||||
MODEL_INPUT_KEYS = (
|
||||
"event_seq",
|
||||
"time_seq",
|
||||
"sex",
|
||||
"padding_mask",
|
||||
"other_type",
|
||||
"other_value",
|
||||
"other_value_kind",
|
||||
"other_time",
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Train DeepHealth with next-token/point supervision")
|
||||
|
||||
parser.add_argument("--data_prefix", type=str, default="ukb")
|
||||
parser.add_argument("--labels_file", type=str, default="labels.csv")
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--extra_info_types_file", type=str, default=None)
|
||||
parser.add_argument("--no_event_interval_years", type=float, default=5.0)
|
||||
parser.add_argument("--include_no_event_in_uts_target", action="store_true")
|
||||
|
||||
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("--n_embd", type=int, default=120)
|
||||
parser.add_argument("--n_head", type=int, default=10)
|
||||
parser.add_argument("--n_hist_layer", type=int, default=12)
|
||||
parser.add_argument("--n_tab_layer", type=int, default=4)
|
||||
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("--dropout", type=float, default=0.0)
|
||||
|
||||
parser.add_argument("--target_mode", type=str, default="uts",
|
||||
choices=["delphi2m", "uts"])
|
||||
parser.add_argument("--readout_name", type=str, default=None,
|
||||
choices=["token", "same_time_group_end", "last_valid"])
|
||||
parser.add_argument("--readout_reduce", type=str, default="mean",
|
||||
choices=["mean", "sum"])
|
||||
parser.add_argument("--t_min", type=float, default=0.0027378507871321013)
|
||||
parser.add_argument("--max_exp_input", type=float, default=60.0)
|
||||
parser.add_argument("--ce_weight", type=float, default=1.0)
|
||||
parser.add_argument("--time_weight", type=float, default=1.0)
|
||||
parser.add_argument("--ignore_no_event_in_delphi2m", action="store_true")
|
||||
|
||||
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()
|
||||
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.target_mode == "uts":
|
||||
args.readout_name = args.readout_name or "same_time_group_end"
|
||||
args.include_no_event_in_uts_target = True
|
||||
else:
|
||||
args.readout_name = args.readout_name or "token"
|
||||
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
|
||||
)
|
||||
return args
|
||||
|
||||
|
||||
def get_lr(epoch: int, args: argparse.Namespace, adaptive_lr: float) -> float:
|
||||
if epoch < args.warmup_epochs:
|
||||
return adaptive_lr * (epoch + 1) / args.warmup_epochs
|
||||
progress = (epoch - args.warmup_epochs) / max(1, args.max_epochs - args.warmup_epochs)
|
||||
cosine = 0.5 * (1 + math.cos(math.pi * progress))
|
||||
return adaptive_lr * (args.min_lr_ratio + cosine * (1 - args.min_lr_ratio))
|
||||
|
||||
|
||||
def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]:
|
||||
non_blocking = device.type == "cuda"
|
||||
return {
|
||||
key: value.to(device, non_blocking=non_blocking)
|
||||
if isinstance(value, torch.Tensor)
|
||||
else value
|
||||
for key, value in batch.items()
|
||||
}
|
||||
|
||||
|
||||
def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
|
||||
return DeepHealth(
|
||||
vocab_size=dataset.vocab_size,
|
||||
n_embd=args.n_embd,
|
||||
n_head=args.n_head,
|
||||
n_hist_layer=args.n_hist_layer,
|
||||
n_tab_layer=args.n_tab_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,
|
||||
extra_pool_reduce=args.extra_pool_reduce,
|
||||
target_mode="next_token",
|
||||
time_mode=args.time_mode,
|
||||
dist_mode="exponential",
|
||||
dropout=args.dropout,
|
||||
)
|
||||
|
||||
|
||||
def build_next_step_readout(args: argparse.Namespace):
|
||||
if args.readout_name == "same_time_group_end":
|
||||
return build_readout("same_time_group_end", reduce=args.readout_reduce)
|
||||
return build_readout(args.readout_name)
|
||||
|
||||
|
||||
def build_next_step_loss(args: argparse.Namespace):
|
||||
if args.target_mode == "delphi2m":
|
||||
ignored_tokens = {PAD_IDX, CHECKUP_IDX}
|
||||
if args.ignore_no_event_in_delphi2m:
|
||||
ignored_tokens.add(NO_EVENT_IDX)
|
||||
return build_loss(
|
||||
"delphi2m",
|
||||
ignored_tokens=ignored_tokens,
|
||||
t_min=args.t_min,
|
||||
max_exp_input=args.max_exp_input,
|
||||
ce_weight=args.ce_weight,
|
||||
time_weight=args.time_weight,
|
||||
)
|
||||
return build_loss(
|
||||
"uts",
|
||||
ignored_idx={PAD_IDX, CHECKUP_IDX},
|
||||
t_min=args.t_min,
|
||||
max_exp_input=args.max_exp_input,
|
||||
)
|
||||
|
||||
|
||||
def build_augmented_next_step_targets(
|
||||
batch_cpu: Dict[str, torch.Tensor],
|
||||
model_out: DeepHealthOutput,
|
||||
include_uts_targets: bool,
|
||||
) -> Dict[str, torch.Tensor]:
|
||||
hidden_len = model_out.hidden.size(1)
|
||||
event_len = int(model_out.event_len)
|
||||
extra_len = hidden_len - event_len
|
||||
device = model_out.hidden.device
|
||||
non_blocking = device.type == "cuda"
|
||||
if extra_len <= 0:
|
||||
targets = {
|
||||
"target_event_seq": batch_cpu["target_event_seq"].to(device, non_blocking=non_blocking),
|
||||
"target_time_seq": batch_cpu["target_time_seq"].to(device, non_blocking=non_blocking),
|
||||
"readout_mask": batch_cpu["readout_mask"].to(device, non_blocking=non_blocking),
|
||||
}
|
||||
if include_uts_targets:
|
||||
targets["target_dt_unique"] = batch_cpu["target_dt_unique"].to(
|
||||
device, non_blocking=non_blocking
|
||||
)
|
||||
targets["target_multi_hot"] = batch_cpu["target_multi_hot"].to(
|
||||
device, non_blocking=non_blocking
|
||||
)
|
||||
return targets
|
||||
|
||||
bsz = batch_cpu["target_event_seq"].size(0)
|
||||
vocab_size = (
|
||||
batch_cpu["target_multi_hot"].size(2)
|
||||
if include_uts_targets
|
||||
else None
|
||||
)
|
||||
other_valid = batch_cpu["other_type"] > 0
|
||||
extra_time = batch_cpu["other_time"].new_zeros(bsz, extra_len)
|
||||
extra_mask = torch.zeros(bsz, extra_len, dtype=torch.bool)
|
||||
for b in range(bsz):
|
||||
unique_time = torch.unique(batch_cpu["other_time"][b, other_valid[b]], sorted=True)
|
||||
n_time = min(int(unique_time.numel()), extra_len)
|
||||
if n_time > 0:
|
||||
extra_time[b, :n_time] = unique_time[:n_time]
|
||||
extra_mask[b, :n_time] = True
|
||||
|
||||
target_event_seq = torch.cat(
|
||||
[
|
||||
batch_cpu["target_event_seq"],
|
||||
torch.full(
|
||||
(bsz, extra_len),
|
||||
PAD_IDX,
|
||||
dtype=batch_cpu["target_event_seq"].dtype,
|
||||
),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
target_time_seq = torch.cat(
|
||||
[
|
||||
batch_cpu["target_time_seq"],
|
||||
torch.zeros(
|
||||
bsz,
|
||||
extra_len,
|
||||
dtype=batch_cpu["target_time_seq"].dtype,
|
||||
),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
readout_mask = torch.cat([batch_cpu["readout_mask"], extra_mask], dim=1)
|
||||
target_dt_unique = None
|
||||
target_multi_hot = None
|
||||
if include_uts_targets:
|
||||
target_dt_unique = torch.cat(
|
||||
[
|
||||
batch_cpu["target_dt_unique"],
|
||||
torch.zeros(
|
||||
bsz,
|
||||
extra_len,
|
||||
dtype=batch_cpu["target_dt_unique"].dtype,
|
||||
),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
target_multi_hot = torch.cat(
|
||||
[
|
||||
batch_cpu["target_multi_hot"],
|
||||
torch.zeros(
|
||||
bsz,
|
||||
extra_len,
|
||||
vocab_size,
|
||||
dtype=batch_cpu["target_multi_hot"].dtype,
|
||||
),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
for b in range(bsz):
|
||||
valid_event = batch_cpu["padding_mask"][b].bool()
|
||||
if not valid_event.any():
|
||||
continue
|
||||
n_event = int(valid_event.sum().item())
|
||||
events = torch.cat(
|
||||
[
|
||||
batch_cpu["event_seq"][b, :n_event],
|
||||
batch_cpu["target_event_seq"][b, n_event - 1:n_event],
|
||||
]
|
||||
)
|
||||
times = torch.cat(
|
||||
[
|
||||
batch_cpu["time_seq"][b, :n_event],
|
||||
batch_cpu["target_time_seq"][b, n_event - 1:n_event],
|
||||
]
|
||||
)
|
||||
valid_full = events > PAD_IDX
|
||||
events = events[valid_full]
|
||||
times = times[valid_full]
|
||||
if events.numel() == 0:
|
||||
continue
|
||||
|
||||
for j in range(extra_len):
|
||||
if not bool(extra_mask[b, j]):
|
||||
continue
|
||||
pos = event_len + j
|
||||
t = extra_time[b, j]
|
||||
future = times > t
|
||||
if not future.any():
|
||||
readout_mask[b, pos] = False
|
||||
continue
|
||||
|
||||
first_idx = int(torch.nonzero(future, as_tuple=False)[0].item())
|
||||
next_time = times[first_idx]
|
||||
next_event = events[first_idx]
|
||||
target_event_seq[b, pos] = next_event
|
||||
target_time_seq[b, pos] = next_time
|
||||
|
||||
if not include_uts_targets:
|
||||
continue
|
||||
|
||||
same_next_time = times == next_time
|
||||
next_events = events[same_next_time]
|
||||
valid_next_events = next_events[
|
||||
(next_events > PAD_IDX) & (next_events < vocab_size)
|
||||
].long()
|
||||
if valid_next_events.numel() == 0:
|
||||
readout_mask[b, pos] = False
|
||||
continue
|
||||
target_multi_hot[b, pos, valid_next_events] = True
|
||||
target_dt_unique[b, pos] = next_time - t
|
||||
|
||||
targets = {
|
||||
"target_event_seq": target_event_seq.to(device, non_blocking=non_blocking),
|
||||
"target_time_seq": target_time_seq.to(device, non_blocking=non_blocking),
|
||||
"readout_mask": readout_mask.to(device, non_blocking=non_blocking),
|
||||
}
|
||||
if include_uts_targets:
|
||||
targets["target_dt_unique"] = target_dt_unique.to(device, non_blocking=non_blocking)
|
||||
targets["target_multi_hot"] = target_multi_hot.to(device, non_blocking=non_blocking)
|
||||
return targets
|
||||
|
||||
|
||||
def compute_next_step_loss(
|
||||
args: argparse.Namespace,
|
||||
model: DeepHealth,
|
||||
readout,
|
||||
criterion,
|
||||
batch: Dict[str, torch.Tensor],
|
||||
device: torch.device,
|
||||
) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]:
|
||||
batch_cpu = batch
|
||||
batch = move_batch_to_device(
|
||||
{key: batch_cpu[key] for key in MODEL_INPUT_KEYS},
|
||||
device,
|
||||
)
|
||||
model_out = model(
|
||||
event_seq=batch["event_seq"],
|
||||
time_seq=batch["time_seq"],
|
||||
sex=batch["sex"],
|
||||
padding_mask=batch["padding_mask"],
|
||||
other_type=batch["other_type"],
|
||||
other_value=batch["other_value"],
|
||||
other_value_kind=batch["other_value_kind"],
|
||||
other_time=batch["other_time"],
|
||||
target_mode="next_token",
|
||||
return_output=True,
|
||||
)
|
||||
if not isinstance(model_out, DeepHealthOutput):
|
||||
raise TypeError("DeepHealth return_output=True must return DeepHealthOutput")
|
||||
targets = build_augmented_next_step_targets(
|
||||
batch_cpu=batch_cpu,
|
||||
model_out=model_out,
|
||||
include_uts_targets=args.target_mode == "uts",
|
||||
)
|
||||
readout_out = readout(
|
||||
hidden=model_out.hidden,
|
||||
time_seq=model_out.time_seq,
|
||||
padding_mask=model_out.padding_mask,
|
||||
readout_mask=targets["readout_mask"]
|
||||
if args.readout_name == "same_time_group_end"
|
||||
else None,
|
||||
)
|
||||
logits = model.calc_risk(readout_out.hidden)
|
||||
|
||||
if args.target_mode == "delphi2m":
|
||||
loss, parts = criterion(
|
||||
logits=logits,
|
||||
target_events=targets["target_event_seq"],
|
||||
target_times=targets["target_time_seq"],
|
||||
current_times=model_out.time_seq,
|
||||
padding_mask=readout_out.readout_mask,
|
||||
return_components=True,
|
||||
)
|
||||
else:
|
||||
loss, parts = criterion(
|
||||
logits=logits,
|
||||
target_multi_hot=targets["target_multi_hot"],
|
||||
target_dt_unique=targets["target_dt_unique"],
|
||||
readout_mask=readout_out.readout_mask,
|
||||
return_components=True,
|
||||
)
|
||||
if not torch.isfinite(loss):
|
||||
raise RuntimeError(f"Loss is not finite: {float(loss.detach().cpu())}")
|
||||
return loss, parts
|
||||
|
||||
|
||||
def run_epoch(
|
||||
logger: logging.Logger,
|
||||
args: argparse.Namespace,
|
||||
model: DeepHealth,
|
||||
readout,
|
||||
criterion,
|
||||
loader: DataLoader,
|
||||
optimizer: AdamW | None,
|
||||
device: torch.device,
|
||||
is_train: bool,
|
||||
) -> float:
|
||||
model.train(is_train)
|
||||
readout.train(is_train)
|
||||
total = torch.zeros((), device=device)
|
||||
n_batches = 0
|
||||
skipped = 0
|
||||
parts_sum: Dict[str, torch.Tensor] = {}
|
||||
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, parts = compute_next_step_loss(args, model, readout, 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
|
||||
for name, value in parts.items():
|
||||
parts_sum[name] = parts_sum.get(name, torch.zeros((), device=device)) + value.detach()
|
||||
if (batch_idx + 1) % progress_interval == 0:
|
||||
avg = total / max(1, n_batches)
|
||||
postfix = {
|
||||
"loss": f"{float(loss.detach().cpu()):.4f}",
|
||||
"avg": f"{float(avg.detach().cpu()):.4f}",
|
||||
"skipped": skipped,
|
||||
}
|
||||
for name, value in parts_sum.items():
|
||||
postfix[name] = f"{float((value / max(1, n_batches)).detach().cpu()):.4f}"
|
||||
progress.set_postfix(postfix)
|
||||
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: HealthDataset,
|
||||
run_name: str,
|
||||
train_subset,
|
||||
val_subset,
|
||||
test_subset,
|
||||
) -> Dict[str, Any]:
|
||||
return {
|
||||
"run_name": run_name,
|
||||
"dataset_class": "NextStepHealthDataset",
|
||||
"collate_fn": "next_step_collate_fn",
|
||||
"model_class": "DeepHealth",
|
||||
"model_target_mode": "next_token",
|
||||
"target_mode": args.target_mode,
|
||||
"dist_mode": "exponential",
|
||||
"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],
|
||||
"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],
|
||||
},
|
||||
"split_sizes": {
|
||||
"train": int(len(train_subset)),
|
||||
"val": int(len(val_subset)),
|
||||
"test": int(len(test_subset)),
|
||||
},
|
||||
"resolved_readout_name": args.readout_name,
|
||||
"resolved_loss_name": args.target_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: (
|
||||
f"{args.time_mode}_exponential_next_token_{args.target_mode}_"
|
||||
f"gap_{args.no_event_interval_years:g}y_{timestamp}"
|
||||
)
|
||||
)
|
||||
logger = setup_logging(run_dir)
|
||||
|
||||
logger.info(f"Starting next-step training run: {run_name}")
|
||||
logger.info(f"Device: {device}")
|
||||
logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}")
|
||||
logger.info(f"readout={args.readout_name}, target_mode={args.target_mode}")
|
||||
|
||||
dataset = HealthDataset(
|
||||
data_prefix=args.data_prefix,
|
||||
labels_file=args.labels_file,
|
||||
no_event_interval_years=args.no_event_interval_years,
|
||||
include_no_event_in_uts_target=args.include_no_event_in_uts_target,
|
||||
extra_info_types=args.extra_info_types,
|
||||
)
|
||||
if args.train_eid_file and args.val_eid_file and args.test_eid_file:
|
||||
train_subset, val_subset, test_subset = split_dataset_by_eid_files(
|
||||
dataset=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_dataset(
|
||||
dataset=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"Samples: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
|
||||
)
|
||||
|
||||
train_loader = DataLoader(
|
||||
train_subset,
|
||||
batch_size=args.batch_size,
|
||||
sampler=RandomSampler(train_subset, generator=torch.Generator().manual_seed(args.seed)),
|
||||
collate_fn=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=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=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, dataset).to(device)
|
||||
readout = build_next_step_readout(args).to(device)
|
||||
criterion = build_next_step_loss(args)
|
||||
optimizer = AdamW(
|
||||
model.parameters(),
|
||||
lr=args.base_lr,
|
||||
betas=tuple(args.betas),
|
||||
weight_decay=args.weight_decay,
|
||||
)
|
||||
adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)
|
||||
|
||||
save_config(
|
||||
args,
|
||||
run_dir / "train_config.json",
|
||||
extra=build_metadata(args, dataset, run_name, train_subset, val_subset, test_subset),
|
||||
)
|
||||
|
||||
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, readout, criterion, train_loader, optimizer, device, True)
|
||||
with torch.no_grad():
|
||||
val_loss = run_epoch(logger, args, model, readout, 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 next-step test split...")
|
||||
model.load_state_dict(torch.load(best_model_path, map_location=device))
|
||||
with torch.no_grad():
|
||||
test_loss = run_epoch(logger, args, model, readout, 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()
|
||||
Reference in New Issue
Block a user