File size: 7,116 Bytes
ae73c7f | 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 164 165 166 167 168 169 170 171 172 173 174 | """
Demonstration of continual learning on successive scientific domains
using the hierarchical Poincaré model + Replay + EWC.
Uses the exact best hyperparameters from the long Optuna study:
lr=3.82e-4, curvature=0.455, hidden=96, batch_size=8,
pred_steps=4, w_phys=9.6e-4, levels=2
"""
from __future__ import annotations
import os, sys, copy
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch
from torch.utils.data import DataLoader
import numpy as np
from tqdm import tqdm
from src.normalization import FieldNormalizer
from src.synthetic_fields import SyntheticWellLike
from src.model import MultiScaleEncoder, HierarchicalHyperbolicPredictor
from src.physics_losses import combined_physics_loss
from src.continual import ReplayBuffer, DiagonalEWC, hyperbolic_distillation_loss, fit_normalizer_for_domain
from src.config import BEST_HPARAMS as BEST
def collate(batch):
return torch.stack([b["fields"] for b in batch])
def make_domain(seed: int, n_channels: int = 2, n_samples=96, n_steps=14):
"""Slightly different synthetic regimes act as successive scientific domains."""
torch.manual_seed(seed)
return SyntheticWellLike(n_samples=n_samples, n_steps=n_steps, height=32, width=32,
n_channels=n_channels, noise=0.12 + 0.04*(seed%3))
def evaluate(model, norm, ds, device, pred_steps=None):
if pred_steps is None:
pred_steps = model.pred_steps
model.eval()
loader = DataLoader(ds, batch_size=8, collate_fn=collate)
losses = []
with torch.no_grad():
for batch in loader:
B,T,C,H,W = batch.shape
batch = batch.to(device)
flat = norm.transform(batch.view(B*T,C,H,W)).view(B,T,C,H,W)
win = 4
if T < win + pred_steps:
continue
x = flat[:, :win]
tgt = torch.stack([model.encode(flat[:, win+s]) for s in range(pred_steps)], 1)
pred = model(x)
losses.append(model.hyperbolic_loss(pred, tgt).item())
model.train()
return float(np.mean(losses)) if losses else 1e6
def train_domain(model, norm, ds, opt, device, epochs=5, replay: ReplayBuffer=None,
ewc: DiagonalEWC=None, teacher=None, mix_replay=0.4):
loader = DataLoader(ds, batch_size=BEST["batch_size"], shuffle=True, collate_fn=collate)
w_phys = BEST["w_phys"]
ps = model.pred_steps
probe = ds[0]["fields"]
domain_c = int(probe.shape[1])
domain_hw = (int(probe.shape[2]), int(probe.shape[3]))
for ep in range(epochs):
for batch in loader:
B,T,C,H,W = batch.shape
batch = batch.to(device)
if replay is not None and len(replay) > 0 and np.random.rand() < mix_replay:
old = replay.sample(max(1, B//2), channels=domain_c, spatial=domain_hw)
if old is not None:
old = old.to(device)
tmin = min(old.size(1), T)
batch = torch.cat([batch[:, :tmin], old[:, :tmin]], dim=0)
B = batch.size(0)
T = tmin
flat = norm.transform(batch.view(B*T if batch.dim()==5 else B*batch.size(1), C, H, W))
if batch.dim() == 5:
flat = flat.view(B, T, C, H, W)
else:
flat = flat.view(B, -1, C, H, W)
T = flat.size(1)
win = 4
if T < win + ps:
continue
x = flat[:, :win]
with torch.no_grad():
tgt = torch.stack([model.encode(flat[:, win+s]) for s in range(ps)], 1)
pred = model(x)
loss = model.hyperbolic_loss(pred, tgt)
loss = loss + combined_physics_loss(flat[:, :win+ps], w_smooth=w_phys, w_temp=w_phys)
if ewc is not None:
loss = loss + ewc.ewc_loss(model)
if teacher is not None:
with torch.no_grad():
t_lat = teacher.encode(x[:, -1] if x.dim()==5 else x)
s_lat = model.encode(x[:, -1] if x.dim()==5 else x)
loss = loss + 0.1 * hyperbolic_distillation_loss(s_lat, t_lat, model.poincare)
opt.zero_grad()
loss.backward()
if ewc is not None and np.random.rand() < 0.25:
ewc.accumulate_fisher(model, loss)
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
if replay is not None:
for i in range(min(8, len(ds))):
replay.add(ds[i]["fields"])
def main():
device = "cpu"
print("=" * 64)
print("Continual learning demo – Replay + EWC on Poincaré hierarchical model")
print("Using Optuna best HPs:", BEST)
print("=" * 64)
domain_specs = [
{"seed": 11, "n_channels": 2},
{"seed": 22, "n_channels": 5},
{"seed": 33, "n_channels": 2},
]
domains = [make_domain(seed=s["seed"], n_channels=s["n_channels"]) for s in domain_specs]
enc = MultiScaleEncoder(hidden=BEST["hidden"], out_dim=8)
model = HierarchicalHyperbolicPredictor(
enc, c=BEST["curvature"], pred_steps=BEST["pred_steps"], levels=BEST["levels"]
).to(device)
opt = torch.optim.Adam(model.parameters(), lr=BEST["lr"])
replay = ReplayBuffer(capacity=128)
ewc = DiagonalEWC(model, lambda_ewc=500.0)
teacher = None
normalizers = []
history = {f"domain_{i}": [] for i in range(3)}
for d_idx, ds in enumerate(domains):
c = domain_specs[d_idx]["n_channels"]
print(f"\n--- Training domain {d_idx+1}/3 (C={c}) ---")
norm = fit_normalizer_for_domain(ds, max_fit=40)
normalizers.append(norm)
train_domain(model, norm, ds, opt, device, epochs=4,
replay=replay,
ewc=ewc if d_idx > 0 else None,
teacher=teacher, mix_replay=0.0 if d_idx == 0 else 0.35)
print(f" replay buffer by channel count: {replay.counts_by_channels()}")
ewc.finalize_domain(model)
teacher = copy.deepcopy(model).eval()
for p in teacher.parameters():
p.requires_grad = False
for j in range(d_idx + 1):
loss_j = evaluate(model, normalizers[j], domains[j], device) # derives pred_steps from model itself
history[f"domain_{j}"].append(loss_j)
print(f" Eval domain {j+1} (C={domain_specs[j]['n_channels']}) loss: {loss_j:.4f}")
print("\n" + "=" * 64)
print("Retention summary (loss after each successive domain)")
for k, v in history.items():
print(f" {k}: {[round(x,4) for x in v]}")
os.makedirs("logs", exist_ok=True)
path = "logs/poincare8d_continual.pt"
torch.save({
"model": model.state_dict(),
"params": BEST,
"normalizers": [n.state_dict() for n in normalizers],
"domain_specs": domain_specs,
"history": history,
}, path)
print(f"\nSaved {path}")
print("Continual learning module integrated and demonstrated (cross-channel-count).")
if __name__ == "__main__":
main()
|