File size: 3,667 Bytes
28d9c3b | 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 | from __future__ import annotations
import sys
from pathlib import Path
import torch
PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))
sys.path.insert(0, str(Path(__file__).resolve().parent))
from common import ( # noqa: E402
build_laplace_data,
checkpoint_state,
load_config,
project_path,
relative_l2,
resolve_device,
resolve_dtype,
seed_everything,
)
from model.bpinn import ( # noqa: E402
build_model,
laplace1d_loss_components,
weighted_loss,
)
DEFAULT_CONFIG = PROJECT_ROOT / "conf" / "config.yaml"
def main() -> None:
config_path = DEFAULT_CONFIG.resolve()
config = load_config(config_path)
common = config["common"]
training = config["training"]
lbfgs_iters = int(training["lbfgs_iters"])
refinement_lr = float(training["refinement_lr"])
if lbfgs_iters <= 0 or refinement_lr <= 0:
raise ValueError("lbfgs-iters and lr must be positive")
device = resolve_device(str(common["device"]))
dtype = resolve_dtype(str(common["dtype"]))
weight_dir = project_path(common["weight_dir"], PROJECT_ROOT)
input_path = weight_dir / training["checkpoint_name"]
output_path = weight_dir / training["refined_checkpoint_name"]
if not input_path.is_file():
raise FileNotFoundError(f"checkpoint not found: {input_path}")
checkpoint = torch.load(input_path, map_location="cpu", weights_only=True)
state, metadata = checkpoint_state(checkpoint)
model_config = metadata.get("model_config", config["model"])
data_config = metadata.get("data_config", config["data"])
loss_weights = metadata.get("loss_weights", config["loss"])
seed = int(metadata.get("seed", common["seed"]))
seed_everything(seed)
data = build_laplace_data(data_config, seed, device, dtype)
model = build_model(model_config, dtype=dtype).to(device=device, dtype=dtype)
model.load_state_dict(state, strict=True)
def loss_value() -> torch.Tensor:
components = laplace1d_loss_components(
model,
data["x_solution"],
data["u_solution"],
data["x_boundary"],
data["u_boundary"],
data["x_pde"],
)
return weighted_loss(components, loss_weights)
lbfgs = torch.optim.LBFGS(
model.parameters(),
lr=refinement_lr,
max_iter=lbfgs_iters,
max_eval=max(1, 2 * lbfgs_iters),
history_size=50,
line_search_fn="strong_wolfe",
)
def closure() -> torch.Tensor:
lbfgs.zero_grad(set_to_none=True)
loss = loss_value()
if not torch.isfinite(loss):
raise FloatingPointError("BPINN refinement loss became non-finite")
loss.backward()
return loss
print(f"Input checkpoint: {input_path}")
print(f"Device: {device}")
lbfgs.step(closure)
with torch.no_grad():
prediction = model.predict_u(data["x_test"])
error = relative_l2(prediction, data["u_test"])
refined = dict(metadata)
refined.update(
{
"case": "laplace1d",
"architecture": "bpinn",
"model_state": model.state_dict(),
"model_config": model_config,
"data_config": data_config,
"loss_weights": loss_weights,
"seed": seed,
"refinement_iters": lbfgs_iters,
"final_l2": error,
}
)
output_path.parent.mkdir(parents=True, exist_ok=True)
torch.save(refined, output_path)
print(f"Relative L2 after refinement={error:.6e}")
print(f"Saved checkpoint: {output_path}")
if __name__ == "__main__":
main()
|