File size: 10,241 Bytes
8f89e30 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 | """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)
|