File size: 8,923 Bytes
e5a6e71 | 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 156 157 158 159 160 161 162 163 | """Formal 300-epoch FluxNet-LAP pushforward training on SCNet DCU."""
import csv
import json
import random
import time
from pathlib import Path
import numpy as np
import torch
from torch.nn import functional as F
from torch.utils.data import DataLoader
from common import (ROOT, Trajectories, Windows, build_model, checkpoint_config_matches, device, load_config,
output_dir, parser, provenance, rollout_metrics, save_json,
setup, validate)
def train_batch(model, window, optimizer, steps):
"""Eq. (6): first-step MSE + terminal MSE; every feedback is detached."""
optimizer.zero_grad(set_to_none=True)
prediction = model(window[:, 0])
first_loss = F.mse_loss(prediction, window[:, 1])
current = prediction.detach()
# Free the first graph before the terminal graph to reduce memory without
# changing gradients or updating parameters between the two losses.
first_loss.backward()
with torch.no_grad():
for _ in range(1, steps - 1):
current = model(current).detach()
prediction = model(current.detach())
terminal_loss = F.mse_loss(prediction, window[:, steps])
terminal_loss.backward()
if not torch.isfinite(first_loss + terminal_loss):
raise FloatingPointError("Nonfinite training loss")
if not torch.stack([torch.isfinite(parameter.grad).all() for parameter in model.parameters()
if parameter.grad is not None]).all():
raise FloatingPointError("Nonfinite parameter gradient")
optimizer.step()
return float(first_loss.detach()), float(terminal_loss.detach())
def atomic_checkpoint(path, payload):
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_suffix(".tmp")
torch.save(payload, temporary)
temporary.replace(path)
def main():
args_parser = parser(__doc__)
args_parser.add_argument("--resume", type=Path)
args = args_parser.parse_args()
cfg = load_config(args.config)
setup(cfg["seed"])
target_device = device()
out = output_dir(cfg)
settings = cfg["training"]
run_info = provenance()
print("FORMAL_TRAINING " + json.dumps({**run_info, "config": cfg}), flush=True)
train = Trajectories(cfg, "train")
val = Trajectories(cfg, "val")
save_json(out / "data_manifest.json", {"train": train.metadata, "val": val.metadata})
generator = torch.Generator().manual_seed(cfg["seed"])
train_loader = DataLoader(Windows(train, settings["pushforward_steps"]),
batch_size=settings["batch_size"], shuffle=True,
num_workers=settings["num_workers"], pin_memory=True,
persistent_workers=settings["num_workers"] > 0,
generator=generator, drop_last=False)
val_loader = DataLoader(Windows(val, settings["pushforward_steps"]), batch_size=cfg["evaluation"]["batch_size"],
shuffle=False, num_workers=0, pin_memory=True)
model = build_model(cfg, target_device)
optimizer = torch.optim.AdamW(model.parameters(), lr=settings["learning_rate"],
weight_decay=settings["weight_decay"])
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer, factor=settings["scheduler_factor"], patience=settings["scheduler_patience"])
epoch_start, best_loss, best_epoch, history = 1, float("inf"), 0, []
if args.resume:
checkpoint = torch.load(args.resume, map_location="cpu", weights_only=False)
if not checkpoint_config_matches(checkpoint["config"], cfg):
raise ValueError("Resume config differs from checkpoint")
model.load_state_dict(checkpoint["model_state_dict"], strict=True)
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
epoch_start = checkpoint["epoch"] + 1
best_loss, best_epoch = checkpoint["best_val_loss"], checkpoint["best_epoch"]
history = checkpoint["history"]
random.setstate(checkpoint["random_state"])
np.random.set_state(checkpoint["numpy_state"])
torch.set_rng_state(checkpoint["torch_state"])
torch.cuda.set_rng_state_all(checkpoint["cuda_state"])
generator.set_state(checkpoint["loader_state"])
print(f"TRAIN samples={len(train_loader.dataset)} batches/epoch={len(train_loader)} "
f"parameters={sum(p.numel() for p in model.parameters())} "
f"epochs={settings['epochs']} start_epoch={epoch_start}", flush=True)
started = time.monotonic()
long_data = None
for epoch in range(epoch_start, settings["epochs"] + 1):
before = time.monotonic()
model.train()
losses = np.zeros(2)
samples = 0
lr = optimizer.param_groups[0]["lr"]
for batch_index, window in enumerate(train_loader, 1):
window = window.to(target_device, non_blocking=True)
first, terminal = train_batch(model, window, optimizer, settings["pushforward_steps"])
losses += np.array([first, terminal]) * len(window)
samples += len(window)
if batch_index % settings["print_every_batches"] == 0 or batch_index == len(train_loader):
print(f"epoch={epoch:03d}/300 batch={batch_index:03d}/{len(train_loader)} "
f"loss={(losses.sum()/samples):.8e} one_step={losses[0]/samples:.8e} "
f"pushforward5={losses[1]/samples:.8e} lr={lr:.3e}", flush=True)
validation = validate(model, val_loader, target_device)
scheduler.step(validation["mse"])
if validation["mse"] < best_loss:
best_loss, best_epoch = validation["mse"], epoch
row = {"epoch": epoch, "train_loss": float(losses.sum()/samples),
"train_step1_mse": float(losses[0]/samples),
"train_step5_mse": float(losses[1]/samples),
"val_mse": validation["mse"],
"val_step1_mse": validation["step1_mse"],
"val_step5_mse": validation["step5_mse"],
**{f"val_mae_{name}": x for name, x in zip(("h", "mx", "my"), validation["mae"])},
"val_negative_h_ratio": validation["negative_h_ratio"], "lr": lr,
"seconds": time.monotonic() - before}
history.append(row)
print("EPOCH " + json.dumps(row), flush=True)
save_json(out / "training_history.json", history)
with open(out / "training_history.csv", "w", newline="") as stream:
writer = csv.DictWriter(stream, fieldnames=row.keys())
writer.writeheader()
writer.writerows(history)
payload = {"model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(),
"scheduler_state_dict": scheduler.state_dict(), "epoch": epoch,
"best_val_loss": best_loss, "best_epoch": best_epoch, "config": cfg,
"history": history, "provenance": run_info,
"random_state": random.getstate(), "numpy_state": np.random.get_state(),
"torch_state": torch.get_rng_state(), "cuda_state": torch.cuda.get_rng_state_all(),
"loader_state": generator.get_state()}
if best_epoch == epoch:
atomic_checkpoint(ROOT / cfg["paths"]["best_checkpoint"], payload)
print(f"BEST checkpoint epoch={best_epoch} val_mse={best_loss:.8e}", flush=True)
if epoch % settings["checkpoint_every_epochs"] == 0 or epoch == settings["epochs"]:
atomic_checkpoint(ROOT / cfg["paths"]["last_checkpoint"], payload)
save_json(out / "training_status.json", {"state": "running", "completed_epochs": epoch,
"best_epoch": best_epoch, "best_val_loss": best_loss,
"elapsed_seconds": time.monotonic() - started, "provenance": run_info})
if epoch % cfg["evaluation"]["long_test_every_epochs"] == 0:
if long_data is None:
long_data = Trajectories(cfg, "test_long")
diagnostic = rollout_metrics(model, long_data, target_device, cfg["evaluation"]["batch_size"])
save_json(out / f"diagnostic_long_epoch{epoch:03d}.json", diagnostic)
print("LONG_TEST_DIAGNOSTIC " + json.dumps({"epoch": epoch, **diagnostic}), flush=True)
assert len(history) == settings["epochs"] and history[-1]["epoch"] == 300
save_json(out / "training_status.json", {"state": "completed", "completed_epochs": 300,
"best_epoch": best_epoch, "best_val_loss": best_loss,
"elapsed_seconds": time.monotonic() - started, "provenance": run_info,
"formal_training": True, "stopped_early": False})
print(f"TRAINING_COMPLETED epochs=300 best_epoch={best_epoch} best_val_mse={best_loss:.8e}", flush=True)
if __name__ == "__main__":
main()
|