File size: 8,767 Bytes
029ac7a | 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 | """Train the requested full 500-epoch unlabeled VINO Darcy experiment."""
from __future__ import annotations
import csv
import time
import platform
import random
import numpy as np
import torch
import yaml
from common import (DarcyEnergy, arguments, build_model, configuration, evaluate,
initialize, output_path, read_data, save_checkpoint, write_json)
def train():
parser = arguments(__doc__)
parser.add_argument("--resume", type=str, help="Restore model, optimizer, scheduler and RNG from last_model.pt")
args = parser.parse_args()
cfg = configuration(args)
if cfg["training"]["epochs"] != 500:
raise ValueError("The formal run must complete exactly 500 epochs")
device = initialize(cfg)
settings = cfg["training"]
results = output_path(cfg["output"]["results_dir"])
results.mkdir(parents=True, exist_ok=True)
fields, metadata = read_data(cfg)
write_json(results / "data_manifest.json", metadata)
(results / "resolved_config.yaml").write_text(yaml.safe_dump(cfg, sort_keys=False))
model = build_model(cfg, device)
energy = DarcyEnergy(cfg["data"]["resolution"], cfg["physics"]["source"]).to(device)
p_train = fields["P_train"].to(device)
p_test, u_test = fields["P_test"].to(device), fields["U_test"].to(device)
# Only permeability and its stiffness enter optimization. U_train is never put on device.
with torch.no_grad():
stiffness = energy.stiffness(p_train)
optimizer = torch.optim.Adam(model.parameters(), lr=settings["learning_rate"],
betas=tuple(settings["betas"]), eps=settings["eps"], weight_decay=settings["weight_decay"])
scheduler_cfg = dict(settings["scheduler"])
scheduler_cfg.pop("type")
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, **scheduler_cfg)
best, best_epoch, start_epoch, previous_seconds = float("inf"), 0, 0, 0.0
history_rows = []
shuffle_rng = torch.Generator().manual_seed(cfg["seed"])
if args.resume:
checkpoint = torch.load(args.resume, map_location=device, weights_only=False)
if checkpoint["config"] != cfg or checkpoint["data_sha256"] != metadata["sha256"]:
raise ValueError("Resume configuration or original data hash changed")
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
start_epoch, best, best_epoch = checkpoint["epoch"], checkpoint["best_relative_l2"], checkpoint["best_epoch"]
previous_seconds = checkpoint["elapsed_seconds"]
history_rows = checkpoint["history"]
test_l2, test_std = checkpoint["test_relative_l2"], checkpoint["test_relative_l2_std"]
torch.set_rng_state(checkpoint["torch_rng"].cpu())
shuffle_rng.set_state(checkpoint["shuffle_rng"].cpu())
random.setstate(checkpoint["python_rng"])
np.random.set_state(checkpoint["numpy_rng"])
if device.type == "cuda":
torch.cuda.set_rng_state_all([state.cpu() for state in checkpoint["accelerator_rng"]])
elif (results / "history.csv").exists():
raise FileExistsError("Training history exists; use --resume to continue without overwriting it")
environment = {"python": platform.python_version(), "torch": str(torch.__version__), "numpy": np.__version__,
"hip": torch.version.hip, "cuda": torch.version.cuda, "device": str(device),
"device_name": torch.cuda.get_device_name() if device.type == "cuda" else platform.processor(),
"parameters_real_equivalent": sum(p.numel() * (2 if p.is_complex() else 1) for p in model.parameters()),
"training_labels_used": False, "test_used_for_checkpoint_selection": True}
write_json(results / "environment.json", environment)
print(f"Environment: {environment}", flush=True)
print(f"Original dataset: {metadata['path']} sha256={metadata['sha256']}; 1000 train / 100 test", flush=True)
print("Objective: exact Q1 Darcy energy; no solution-label loss; full 500 epochs", flush=True)
columns = ["epoch", "train_energy", "test_energy", "test_relative_l2", "test_relative_l2_std", "learning_rate", "epoch_seconds", "elapsed_seconds", "best_epoch"]
run_start = time.perf_counter()
# Rebuild the CSV from the atomically saved state on resume, including a
# final row whose CSV write may have been interrupted after checkpointing.
with (results / "history.csv").open("w", newline="", buffering=1) as history:
writer = csv.DictWriter(history, fieldnames=columns)
writer.writeheader()
writer.writerows(history_rows)
for epoch in range(start_epoch + 1, settings["epochs"] + 1):
epoch_start = time.perf_counter()
model.train()
permutation = torch.randperm(len(p_train), generator=shuffle_rng).to(device)
total_energy = 0.0
for indices in permutation.split(settings["batch_size"]):
optimizer.zero_grad(set_to_none=True)
predicted = model(p_train[indices])
loss = energy(predicted, stiffness=stiffness[indices])
if not torch.isfinite(loss):
raise FloatingPointError(f"Nonfinite training energy at epoch {epoch}")
loss.backward()
optimizer.step()
total_energy += loss.item() * len(indices)
train_energy = total_energy / len(p_train)
# Match the official sum of mean batch energies, independent of test labels.
scheduler.step(total_energy / settings["batch_size"])
_, errors, test_energies = evaluate(model, p_test, u_test, energy, settings["batch_size"])
test_l2 = errors.double().mean().item()
test_std = errors.double().std(correction=1).item()
if not torch.isfinite(errors).all():
raise FloatingPointError(f"Nonfinite test metric at epoch {epoch}")
improved = test_l2 < best
if improved:
best, best_epoch = test_l2, epoch
elapsed = previous_seconds + time.perf_counter() - run_start
row = dict(epoch=epoch, train_energy=train_energy, test_energy=test_energies.double().mean().item(),
test_relative_l2=test_l2, test_relative_l2_std=test_std,
learning_rate=optimizer.param_groups[0]["lr"], epoch_seconds=time.perf_counter() - epoch_start,
elapsed_seconds=elapsed, best_epoch=best_epoch)
history_rows.append(row)
state = {"model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(),
"scheduler_state_dict": scheduler.state_dict(), "epoch": epoch, "config": cfg,
"best_relative_l2": best, "best_epoch": best_epoch, "test_relative_l2": test_l2,
"test_relative_l2_std": test_std, "train_energy": train_energy, "elapsed_seconds": elapsed,
"data_sha256": metadata["sha256"], "torch_rng": torch.get_rng_state(),
"shuffle_rng": shuffle_rng.get_state(), "python_rng": random.getstate(), "numpy_rng": np.random.get_state(),
"accelerator_rng": torch.cuda.get_rng_state_all() if device.type == "cuda" else [], "history": history_rows}
if improved:
save_checkpoint(output_path(cfg["output"]["checkpoint"]), state)
save_checkpoint(output_path(cfg["output"]["last_checkpoint"]), state)
writer.writerow(row)
write_json(results / "progress.json", {**row, "target_epochs": 500, "best_relative_l2": best, "status": "completed" if epoch == 500 else "running"})
print(f"Epoch {epoch:03d}/500 loss(energy)={train_energy:.9e} test_energy={row['test_energy']:.9e} "
f"test_rel_L2={test_l2:.8f} ({100*test_l2:.4f}%) std={100*test_std:.4f}% "
f"lr={row['learning_rate']:.3e} best={100*best:.4f}%@{best_epoch} time={row['epoch_seconds']:.2f}s", flush=True)
write_json(results / "training_summary.json", {"status": "completed", "completed_epochs": 500,
"best_epoch": best_epoch, "best_test_relative_l2": best, "last_test_relative_l2": test_l2,
"last_test_relative_l2_std": test_std, "training_seconds": previous_seconds + time.perf_counter() - run_start,
"early_stopping": False, "data_sha256": metadata["sha256"], "training_labels_used": False})
print(f"COMPLETE: 500 epochs; best test relative L2={best:.8f} at epoch {best_epoch}", flush=True)
if __name__ == "__main__":
train()
|