Add training scripts for all-future and next-step supervision with DeepHealth
- Implement `train_all_future.py` for training with query-conditioned all-future supervision. - Implement `train_next_step.py` for training with next-token/next-time-point supervision. - Introduce `train_util.py` for shared utility functions including logging, dataset splitting, and model checkpointing. - Enhance argument parsing for both training scripts to accommodate new parameters. - Update loss functions and model configurations to support the new training paradigms.
This commit is contained in:
196
train_util.py
Normal file
196
train_util.py
Normal file
@@ -0,0 +1,196 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterable, Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.optim import AdamW
|
||||
from torch.utils.data import Subset
|
||||
|
||||
from dataset import AllFutureHealthDataset, HealthDataset
|
||||
from models import DeepHealth
|
||||
|
||||
|
||||
def setup_logging(run_dir: Path) -> logging.Logger:
|
||||
run_dir.mkdir(parents=True, exist_ok=True)
|
||||
logger = logging.getLogger("DeepHealth")
|
||||
logger.setLevel(logging.INFO)
|
||||
logger.handlers.clear()
|
||||
|
||||
formatter = logging.Formatter(
|
||||
"%(asctime)s - %(levelname)s - %(message)s",
|
||||
datefmt="%Y-%m-%d %H:%M:%S",
|
||||
)
|
||||
|
||||
console_handler = logging.StreamHandler(sys.stdout)
|
||||
console_handler.setLevel(logging.INFO)
|
||||
console_handler.setFormatter(formatter)
|
||||
logger.addHandler(console_handler)
|
||||
|
||||
file_handler = logging.FileHandler(run_dir / "train.log", mode="w")
|
||||
file_handler.setLevel(logging.INFO)
|
||||
file_handler.setFormatter(formatter)
|
||||
logger.addHandler(file_handler)
|
||||
|
||||
return logger
|
||||
|
||||
|
||||
def set_seed(seed: int) -> None:
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed(seed)
|
||||
|
||||
|
||||
def load_extra_info_types_file(path: str) -> list[int]:
|
||||
file_path = Path(path)
|
||||
if not file_path.is_file():
|
||||
raise FileNotFoundError(f"extra_info_types_file not found: {path}")
|
||||
|
||||
text = file_path.read_text(encoding="utf-8").strip()
|
||||
if not text:
|
||||
return []
|
||||
|
||||
if text.startswith("["):
|
||||
raw_items = json.loads(text)
|
||||
if not isinstance(raw_items, list):
|
||||
raise ValueError("extra_info_types_file JSON must be a list")
|
||||
else:
|
||||
raw_items = []
|
||||
for line in text.splitlines():
|
||||
line = line.split("#", 1)[0].strip()
|
||||
if line:
|
||||
raw_items.extend(line.replace(",", " ").replace(";", " ").split())
|
||||
|
||||
try:
|
||||
return [int(x) for x in raw_items]
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError(f"Invalid extra info type id in {path}") from exc
|
||||
|
||||
|
||||
def configure_torch_for_training(device: torch.device) -> None:
|
||||
if device.type != "cuda":
|
||||
return
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
torch.backends.cudnn.allow_tf32 = True
|
||||
if hasattr(torch, "set_float32_matmul_precision"):
|
||||
torch.set_float32_matmul_precision("high")
|
||||
|
||||
|
||||
def resolve_device(device_arg: str) -> torch.device:
|
||||
requested = device_arg.strip().lower()
|
||||
if requested == "cpu":
|
||||
return torch.device("cpu")
|
||||
if requested == "cuda":
|
||||
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
if requested.startswith("cuda:"):
|
||||
if not torch.cuda.is_available():
|
||||
return torch.device("cpu")
|
||||
index = int(requested.split(":", 1)[1])
|
||||
if index < 0 or index >= torch.cuda.device_count():
|
||||
raise ValueError(f"Requested CUDA device is out of range: {device_arg}")
|
||||
return torch.device(f"cuda:{index}")
|
||||
raise ValueError(f"Unsupported device: {device_arg}")
|
||||
|
||||
|
||||
def split_dataset(
|
||||
dataset: HealthDataset,
|
||||
train_ratio: float,
|
||||
val_ratio: float,
|
||||
test_ratio: float,
|
||||
seed: int,
|
||||
) -> Tuple[Subset, Subset, Subset]:
|
||||
total = train_ratio + val_ratio + test_ratio
|
||||
if not np.isclose(total, 1.0, atol=1e-6):
|
||||
raise ValueError(f"train/val/test ratios must sum to 1.0, got {total}")
|
||||
|
||||
indices = np.random.RandomState(seed).permutation(len(dataset))
|
||||
n_train = int(len(dataset) * train_ratio)
|
||||
n_val = int(len(dataset) * val_ratio)
|
||||
return (
|
||||
Subset(dataset, indices[:n_train]),
|
||||
Subset(dataset, indices[n_train:n_train + n_val]),
|
||||
Subset(dataset, indices[n_train + n_val:]),
|
||||
)
|
||||
|
||||
|
||||
def split_all_future_datasets(
|
||||
train_dataset: AllFutureHealthDataset,
|
||||
val_dataset: AllFutureHealthDataset,
|
||||
test_dataset: AllFutureHealthDataset,
|
||||
train_ratio: float,
|
||||
val_ratio: float,
|
||||
test_ratio: float,
|
||||
seed: int,
|
||||
) -> Tuple[Subset, Subset, Subset]:
|
||||
total = train_ratio + val_ratio + test_ratio
|
||||
if not np.isclose(total, 1.0, atol=1e-6):
|
||||
raise ValueError(f"train/val/test ratios must sum to 1.0, got {total}")
|
||||
|
||||
patient_indices = np.random.RandomState(seed).permutation(len(train_dataset.patients))
|
||||
n_train = int(len(patient_indices) * train_ratio)
|
||||
n_val = int(len(patient_indices) * val_ratio)
|
||||
train_patient_idx = patient_indices[:n_train]
|
||||
val_patient_set = set(int(x) for x in patient_indices[n_train:n_train + n_val])
|
||||
test_patient_set = set(int(x) for x in patient_indices[n_train + n_val:])
|
||||
|
||||
val_query_idx = [
|
||||
i for i, (pidx, _t_query) in enumerate(val_dataset.valid_queries)
|
||||
if int(pidx) in val_patient_set
|
||||
]
|
||||
test_query_idx = [
|
||||
i for i, (pidx, _t_query) in enumerate(test_dataset.valid_queries)
|
||||
if int(pidx) in test_patient_set
|
||||
]
|
||||
if not val_query_idx:
|
||||
raise ValueError("All-future validation split has no valid query samples.")
|
||||
if not test_query_idx:
|
||||
raise ValueError("All-future test split has no valid query samples.")
|
||||
|
||||
return (
|
||||
Subset(train_dataset, train_patient_idx),
|
||||
Subset(val_dataset, np.asarray(val_query_idx, dtype=np.int64)),
|
||||
Subset(test_dataset, np.asarray(test_query_idx, dtype=np.int64)),
|
||||
)
|
||||
|
||||
|
||||
def build_optimizer(args: Any, model: DeepHealth) -> AdamW:
|
||||
return AdamW(
|
||||
model.parameters(),
|
||||
lr=args.base_lr,
|
||||
betas=tuple(args.betas),
|
||||
weight_decay=args.weight_decay,
|
||||
)
|
||||
|
||||
|
||||
def set_optimizer_lr(optimizer: AdamW, lr: float) -> None:
|
||||
for param_group in optimizer.param_groups:
|
||||
param_group["lr"] = lr
|
||||
|
||||
|
||||
def save_checkpoint(model: DeepHealth, checkpoint_path: Path) -> None:
|
||||
torch.save(model.state_dict(), checkpoint_path)
|
||||
|
||||
|
||||
def save_config(
|
||||
args: Any,
|
||||
config_path: Path,
|
||||
extra: Dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
config: Dict[str, Any] = {}
|
||||
for key, value in vars(args).items():
|
||||
if isinstance(value, tuple):
|
||||
config[key] = list(value)
|
||||
elif isinstance(value, list):
|
||||
config[key] = value
|
||||
elif isinstance(value, (int, float, str, bool, type(None))):
|
||||
config[key] = value
|
||||
else:
|
||||
config[key] = str(value)
|
||||
if extra:
|
||||
config.update(extra)
|
||||
config_path.write_text(json.dumps(config, indent=2), encoding="utf-8")
|
||||
Reference in New Issue
Block a user