| """Full paper training: python scripts/train.py --config config/config.yaml.""" |
|
|
| import argparse |
| import csv |
| import os |
| import platform |
| import sys |
| import time |
| from pathlib import Path |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
| import torch |
| from torch.utils.data import DataLoader, TensorDataset |
| from model.lino import LiNO |
| from common import (atomic_checkpoint, compute_device, decode, encode, file_digest, |
| fit_normalizers, load_checkpoint, load_config, make_inputs, |
| package_path, read_fields, relative_l2, save_json, |
| seed_everything, split_indices) |
|
|
|
|
| @torch.no_grad() |
| def evaluate(model, loader, stats, device): |
| model.eval() |
| total = torch.zeros(4, device=device) |
| count = 0 |
| for inputs, targets in loader: |
| inputs, targets = inputs.to(device), targets.to(device) |
| prediction = model(inputs) |
| mse = (prediction - targets).square().mean() |
| rel = relative_l2(prediction, targets).mean() |
| physical_prediction, physical_target = decode(prediction, stats), decode(targets, stats) |
| physical_rel = relative_l2(physical_prediction, physical_target).mean() |
| physical_mse = (physical_prediction - physical_target).square().mean() |
| total += torch.stack([mse + rel, mse, physical_rel, physical_mse]) * len(inputs) |
| count += len(inputs) |
| values = (total / count).cpu().tolist() |
| if not all(torch.isfinite(torch.tensor(values))): |
| raise FloatingPointError("Nonfinite validation metrics") |
| return dict(zip(["val_loss", "val_mse_normalized", "val_relative_l2", "val_mse_physical"], values)) |
|
|
|
|
| def train(config, resume=None): |
| seed_everything(config["seed"]) |
| settings = config["training"] |
| device = compute_device(settings["device"]) |
| output = package_path(config["results"]["directory"]) |
| output.mkdir(parents=True, exist_ok=True) |
| initial, final = read_fields(config) |
| train_ids, val_ids = split_indices(config) |
| stats = fit_normalizers(initial, final, train_ids, config["data"]["normalization_eps"]) |
| restored = load_checkpoint(resume) if resume else None |
| if restored: |
| if restored["architecture"] != config["model"]: |
| raise ValueError("Resume architecture differs from config") |
| train_ids, val_ids, stats = restored["train_indices"], restored["val_indices"], restored["normalizers"] |
| inputs, targets = make_inputs(initial, stats), encode(final, stats, "u") |
| loader_args = dict(batch_size=settings["batch_size"], num_workers=settings["num_workers"], pin_memory=True, drop_last=False) |
| training = DataLoader(TensorDataset(inputs[train_ids], targets[train_ids]), shuffle=True, **loader_args) |
| validation = DataLoader(TensorDataset(inputs[val_ids], targets[val_ids]), shuffle=False, **loader_args) |
| model = LiNO(**config["model"]).to(device) |
| optimizer = torch.optim.AdamW(model.parameters(), lr=settings["learning_rate"], |
| weight_decay=settings["weight_decay"], betas=tuple(settings["betas"]), eps=settings["optimizer_eps"]) |
| scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=settings["scheduler_step"], gamma=settings["scheduler_gamma"]) |
| best, best_epoch, start_epoch = float("inf"), 0, 0 |
| if restored: |
| model.load_state_dict(restored["model_state"], strict=True) |
| optimizer.load_state_dict(restored["optimizer_state"]) |
| scheduler.load_state_dict(restored["scheduler_state"]) |
| best, best_epoch, start_epoch = restored["best_val_relative_l2"], restored["best_epoch"], restored["epoch"] |
| torch.set_rng_state(restored["torch_rng_state"]) |
| torch.cuda.set_rng_state_all(restored["accelerator_rng_states"]) |
| data_path = package_path(config["data"]["directory"]) / config["data"]["filename"] |
| parameter_count = sum(p.numel() for p in model.parameters()) |
| manifest = dict(config=config, architecture=model.architecture, parameter_count=parameter_count, |
| dataset_path=str(data_path), dataset_sha256=file_digest(data_path), |
| raw_shape=config["data"]["expected_shape"], resolution=config["data"]["resolution"], |
| n_train=len(train_ids), n_val=len(val_ids), seed=config["seed"], |
| dtype="float32", optimizer="AdamW", precision="float32", early_stopping=False, |
| torch_version=str(torch.__version__), hip_version=torch.version.hip, |
| python_version=platform.python_version(), hostname=platform.node(), |
| slurm_job_id=os.environ.get("SLURM_JOB_ID"), device=torch.cuda.get_device_name(device), |
| source_files={str(p.relative_to(Path(__file__).resolve().parents[1])):file_digest(p) |
| for folder in ["model", "scripts"] for p in (Path(__file__).resolve().parents[1]/folder).glob("*.py")}) |
| save_json(output / "run_manifest.json", manifest) |
| save_json(output / "split_indices.json", {"seed":config["seed"], "train":train_ids.tolist(), "validation":val_ids.tolist()}) |
| history_path = output / "history.csv" |
| columns = ["epoch", "train_loss", "train_mse_normalized", "train_relative_l2_normalized", "val_loss", "val_mse_normalized", "val_relative_l2", "val_mse_physical", "lr_used", "lr_next", "epoch_seconds", "best_val_relative_l2"] |
| print(f"FORMAL RUN device={manifest['device']} job={manifest['slurm_job_id']} params={parameter_count} train={len(train_ids)} val={len(val_ids)} resolution={inputs.shape[1]} epochs={settings['epochs']} batch={settings['batch_size']}", flush=True) |
| started = time.monotonic() |
| with open(history_path, "a" if resume else "w", newline="", buffering=1) as stream: |
| writer = csv.DictWriter(stream, fieldnames=columns) |
| if not resume: |
| writer.writeheader() |
| for epoch in range(start_epoch + 1, settings["epochs"] + 1): |
| epoch_start = time.monotonic() |
| lr_used = optimizer.param_groups[0]["lr"] |
| model.train() |
| totals = torch.zeros(3, device=device) |
| count = 0 |
| for step, (batch_x, batch_y) in enumerate(training, 1): |
| batch_x, batch_y = batch_x.to(device), batch_y.to(device) |
| optimizer.zero_grad(set_to_none=True) |
| prediction = model(batch_x) |
| mse = (prediction - batch_y).square().mean() |
| relative = relative_l2(prediction, batch_y).mean() |
| loss = mse + relative |
| loss.backward() |
| gradient_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), settings["grad_clip"], error_if_nonfinite=True) |
| optimizer.step() |
| totals += torch.stack([loss.detach(), mse.detach(), relative.detach()]) * len(batch_x) |
| count += len(batch_x) |
| if step % settings["log_interval"] == 0 or step == len(training): |
| print(f"epoch={epoch:03d}/{settings['epochs']} step={step:03d}/{len(training)} loss={loss.item():.7e} mse={mse.item():.7e} rel_l2_norm={relative.item():.7e} grad_norm={gradient_norm.item():.5e}", flush=True) |
| scheduler.step() |
| metrics = evaluate(model, validation, stats, device) |
| train_loss, train_mse, train_rel = (totals / count).cpu().tolist() |
| improved = metrics["val_relative_l2"] < best |
| if improved: |
| best, best_epoch = metrics["val_relative_l2"], epoch |
| row = dict(epoch=epoch, train_loss=train_loss, train_mse_normalized=train_mse, |
| train_relative_l2_normalized=train_rel, **metrics, lr_used=lr_used, |
| lr_next=optimizer.param_groups[0]["lr"], epoch_seconds=time.monotonic()-epoch_start, |
| best_val_relative_l2=best) |
| state = dict(format="lino-burgers-v1", architecture=model.architecture, |
| model_state=model.state_dict(), optimizer_state=optimizer.state_dict(), |
| scheduler_state=scheduler.state_dict(), epoch=epoch, best_epoch=best_epoch, |
| best_val_relative_l2=best, normalizers=stats, train_indices=train_ids, |
| val_indices=val_ids, data_contract={"resolution":config["data"]["resolution"], |
| "interpolation":config["data"]["interpolation"], "input_key":config["data"]["input_key"], |
| "target_key":config["data"]["target_key"], "dataset_sha256":manifest["dataset_sha256"]}, |
| torch_rng_state=torch.get_rng_state(), accelerator_rng_states=torch.cuda.get_rng_state_all(), |
| validation_metrics=metrics) |
| if improved: |
| atomic_checkpoint(settings["checkpoint"], state) |
| atomic_checkpoint(settings["last_checkpoint"], state) |
| writer.writerow(row) |
| save_json(output / "training_status.json", {**row, "best_epoch":best_epoch, "status":"running", "completed_epochs":epoch}) |
| print(f"EPOCH {epoch:03d}/{settings['epochs']} train_loss={train_loss:.8e} val_loss={metrics['val_loss']:.8e} val_relative_l2={metrics['val_relative_l2']:.8e} best={best:.8e} best_epoch={best_epoch} lr={lr_used:.6e} seconds={row['epoch_seconds']:.2f}", flush=True) |
| summary = dict(status="completed", completed_epochs=settings["epochs"], best_epoch=best_epoch, |
| best_val_relative_l2=best, paper_relative_l2=0.0109, |
| ratio_to_paper=best/0.0109, training_seconds=time.monotonic()-started, |
| n_train=len(train_ids), n_val=len(val_ids), seed=config["seed"], |
| checkpoint=str(settings["checkpoint"]), slurm_job_id=manifest["slurm_job_id"]) |
| save_json(output / "training_summary.json", summary) |
| save_json(output / "training_status.json", summary) |
| print("TRAINING COMPLETED", summary, flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--config", default="config/config.yaml") |
| parser.add_argument("--resume", default=None, help="Restore a full last checkpoint if interrupted") |
| args = parser.parse_args() |
| train(load_config(args.config), args.resume) |
|
|
