1041 lines
37 KiB
Python
1041 lines
37 KiB
Python
# train_ddp_fno_encoderfno_ddpsafe_split_save_all_fields.py
|
|
# =========================================================
|
|
# ✅ DDP-safe split + training
|
|
# ✅ After training:
|
|
# - Load best model
|
|
# - Infer ALL train/test samples (DDP)
|
|
# - Save per-sample errors + ALL E_pred/E_true + coords to MAT
|
|
# - Also save an example sample MAT + PNG plot (optional)
|
|
#
|
|
# ✅ Adapted for NEW dataset: metalens_dataset.mat (NO standardization / NO denorm)
|
|
# - Rimg: (N,1,512,512)
|
|
# - coords: (M,2) in [-1,1]
|
|
# - Ere_flat: (N,M) normalized
|
|
# - Eim_flat: (N,M) normalized
|
|
# =========================================================
|
|
|
|
import os
|
|
import math
|
|
import datetime
|
|
import numpy as np
|
|
import torch
|
|
from tqdm.auto import tqdm
|
|
from scipy.io import savemat
|
|
import torch.nn as nn
|
|
from torch.utils.data import DataLoader, Subset, ConcatDataset
|
|
import matplotlib.pyplot as plt
|
|
from scipy.interpolate import griddata
|
|
from pathlib import Path
|
|
|
|
import torch.distributed as dist
|
|
from torch.nn.parallel import DistributedDataParallel as DDP
|
|
from torch.utils.data.distributed import DistributedSampler
|
|
|
|
# Existing API (must be updated to read metalens_dataset.mat)
|
|
from getdata import build_datasets
|
|
|
|
|
|
# =========================================================
|
|
# DDP helpers
|
|
# =========================================================
|
|
def ddp_setup():
|
|
dist.init_process_group(backend="nccl")
|
|
local_rank = int(os.environ["LOCAL_RANK"])
|
|
torch.cuda.set_device(local_rank)
|
|
device = torch.device(f"cuda:{local_rank}")
|
|
return device, local_rank
|
|
|
|
|
|
def ddp_cleanup():
|
|
if dist.is_available() and dist.is_initialized():
|
|
dist.destroy_process_group()
|
|
|
|
|
|
def is_main_process():
|
|
return (not dist.is_available()) or (not dist.is_initialized()) or dist.get_rank() == 0
|
|
|
|
|
|
def world_size():
|
|
return dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1
|
|
|
|
|
|
def broadcast_tensor_(t: torch.Tensor, src: int = 0):
|
|
if world_size() > 1:
|
|
dist.broadcast(t, src=src)
|
|
return t
|
|
|
|
|
|
def all_gather_object(obj):
|
|
"""Gather python objects to every rank."""
|
|
if world_size() == 1:
|
|
return [obj]
|
|
out = [None for _ in range(world_size())]
|
|
dist.all_gather_object(out, obj)
|
|
return out
|
|
|
|
|
|
# =========================================================
|
|
# Norm
|
|
# =========================================================
|
|
def make_gn(num_channels: int, max_groups: int = 8) -> nn.GroupNorm:
|
|
for g in [max_groups, 4, 2, 1]:
|
|
if num_channels % g == 0:
|
|
return nn.GroupNorm(num_groups=g, num_channels=num_channels)
|
|
return nn.GroupNorm(num_groups=1, num_channels=num_channels)
|
|
|
|
|
|
# =========================================================
|
|
# Loss: Relative L2 (on flattened M points)
|
|
# =========================================================
|
|
class RelL2Loss(nn.Module):
|
|
def __init__(self, eps=1e-12):
|
|
super().__init__()
|
|
self.eps = eps
|
|
|
|
def forward(self, pred: torch.Tensor, true: torch.Tensor) -> torch.Tensor:
|
|
# pred/true: (B,M,2)
|
|
num = torch.sum((pred - true) ** 2, dim=(1, 2))
|
|
den = torch.sum(true ** 2, dim=(1, 2)) + self.eps
|
|
return torch.mean(num / den)
|
|
|
|
|
|
def rel_l2_per_sample(pred: torch.Tensor, true: torch.Tensor, eps=1e-12) -> torch.Tensor:
|
|
num = torch.sum((pred - true) ** 2, dim=(1, 2))
|
|
den = torch.sum(true ** 2, dim=(1, 2)) + eps
|
|
return num / den # (B,)
|
|
|
|
|
|
# =========================================================
|
|
# FNO core (self-contained) -- FIXED FFTConv2d (modes clamp)
|
|
# =========================================================
|
|
class FFTConv2d(nn.Module):
|
|
def __init__(self, in_channels: int, out_channels: int, modes: int):
|
|
super().__init__()
|
|
self.out_channels = out_channels
|
|
self.modes = int(modes)
|
|
self.w = nn.Parameter(
|
|
torch.rand(in_channels, out_channels, 2 * self.modes, self.modes, 2)
|
|
/ (in_channels * out_channels)
|
|
)
|
|
|
|
@staticmethod
|
|
def cmul2d(a, b):
|
|
return torch.einsum("bixy,ioxy->boxy", a, b)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
B, C, H, W = x.shape
|
|
device = x.device
|
|
|
|
x_ft = torch.fft.rfft2(x) # (B,C,H,Wf)
|
|
Wf = x_ft.size(-1)
|
|
|
|
my = min(self.modes, H // 2)
|
|
mx = min(self.modes, Wf)
|
|
|
|
out_ft = torch.zeros(B, self.out_channels, H, Wf, dtype=torch.cfloat, device=device)
|
|
|
|
x_ft = torch.fft.fftshift(x_ft, dim=-2)
|
|
cy = H // 2
|
|
|
|
a = x_ft[..., cy - my: cy + my, :mx] # (B,inC,2*my,mx)
|
|
w_c = torch.view_as_complex(self.w) # (inC,outC,2*modes,modes)
|
|
b = w_c[..., self.modes - my: self.modes + my, :mx]
|
|
out_ft[..., cy - my: cy + my, :mx] = self.cmul2d(a, b)
|
|
|
|
out_ft = torch.fft.ifftshift(out_ft, dim=-2)
|
|
x = torch.fft.irfft2(out_ft, s=(H, W))
|
|
return x
|
|
|
|
|
|
class FNOBlock2d(nn.Module):
|
|
def __init__(self, modes: int, width: int, norm: str = "gn", dropout: float = 0.10, layerscale=0.10):
|
|
super().__init__()
|
|
self.fftconv = FFTConv2d(width, width, modes)
|
|
self.conv = nn.Conv2d(width, width, 1, bias=False)
|
|
|
|
self.norm = nn.BatchNorm2d(width) if norm == "bn" else make_gn(width)
|
|
self.act = nn.GELU()
|
|
self.drop = nn.Dropout2d(p=float(dropout)) if dropout and dropout > 0 else nn.Identity()
|
|
self.gamma = nn.Parameter(layerscale * torch.ones(1, width, 1, 1)) if layerscale and layerscale > 0 else None
|
|
|
|
def forward(self, x):
|
|
y = self.fftconv(x) + self.conv(x)
|
|
x = x + (self.gamma * y if self.gamma is not None else y)
|
|
x = self.norm(x)
|
|
x = self.act(x)
|
|
x = self.drop(x)
|
|
return x
|
|
|
|
|
|
class FNOModel2d(nn.Module):
|
|
def __init__(
|
|
self,
|
|
modes=6,
|
|
width=24,
|
|
blocks=4,
|
|
padding=0,
|
|
in_channels=3,
|
|
out_channels=2,
|
|
norm="gn",
|
|
dropout=0.10,
|
|
layerscale=0.10,
|
|
):
|
|
super().__init__()
|
|
self.conv_in = nn.Conv2d(in_channels, width, 1, bias=True)
|
|
|
|
self.padding = int(padding)
|
|
self.pad_in = nn.ConstantPad2d(self.padding, 0.0) if self.padding > 0 else nn.Identity()
|
|
self.pad_out = nn.ConstantPad2d(-self.padding, 0.0) if self.padding > 0 else nn.Identity()
|
|
|
|
self.fno_blocks = nn.Sequential(
|
|
*[FNOBlock2d(modes, width, norm=norm, dropout=dropout, layerscale=layerscale) for _ in range(blocks)]
|
|
)
|
|
|
|
self.conv_out = nn.Sequential(
|
|
nn.Conv2d(width, width, 1),
|
|
nn.GELU(),
|
|
nn.Dropout2d(dropout),
|
|
nn.Conv2d(width, out_channels, 1),
|
|
)
|
|
|
|
def forward(self, x):
|
|
x = self.conv_in(x)
|
|
x = self.pad_in(x)
|
|
x = self.fno_blocks(x)
|
|
x = self.pad_out(x)
|
|
x = self.conv_out(x)
|
|
return x
|
|
|
|
|
|
# =========================================================
|
|
# Anti-alias Geometry Encoder
|
|
# =========================================================
|
|
class GeomEncoder(nn.Module):
|
|
def __init__(self, out_ch=8, drop=0.10):
|
|
super().__init__()
|
|
self.prepool = nn.AvgPool2d(kernel_size=6, stride=6) # 512 -> ~85
|
|
|
|
self.net = nn.Sequential(
|
|
nn.Conv2d(1, 16, 3, padding=1, bias=False),
|
|
make_gn(16), nn.GELU(), nn.Dropout2d(drop),
|
|
|
|
nn.Conv2d(16, 24, 3, stride=2, padding=1, bias=False), # ~85 -> ~43
|
|
make_gn(24), nn.GELU(), nn.Dropout2d(drop),
|
|
|
|
nn.Conv2d(24, 32, 3, stride=2, padding=1, bias=False), # ~43 -> ~22
|
|
make_gn(32), nn.GELU(), nn.Dropout2d(drop),
|
|
|
|
nn.Conv2d(32, out_ch, 1, bias=True),
|
|
make_gn(out_ch), nn.GELU(),
|
|
)
|
|
self.pool_to = None
|
|
|
|
def forward(self, x_img, Ny, Nx):
|
|
h = self.prepool(x_img)
|
|
h = self.net(h)
|
|
if (self.pool_to is None) or (self.pool_to.output_size != (Ny, Nx)):
|
|
self.pool_to = nn.AdaptiveAvgPool2d((Ny, Nx))
|
|
return self.pool_to(h)
|
|
|
|
|
|
# =========================================================
|
|
# Optional Fourier features (DEFAULT OFF)
|
|
# =========================================================
|
|
class FourierFeatures2D(nn.Module):
|
|
def __init__(self, K=2, scale=1.0):
|
|
super().__init__()
|
|
self.K = int(K)
|
|
self.scale = float(scale)
|
|
freqs = torch.linspace(1.0, self.K, steps=self.K) * self.scale
|
|
self.register_buffer("freqs", freqs, persistent=False)
|
|
|
|
def forward(self, xy):
|
|
x = xy[:, 0:1]
|
|
y = xy[:, 1:2]
|
|
f = self.freqs.view(1, self.K, 1, 1)
|
|
x_proj = 2 * math.pi * x * f
|
|
y_proj = 2 * math.pi * y * f
|
|
return torch.cat(
|
|
[torch.sin(x_proj), torch.cos(x_proj), torch.sin(y_proj), torch.cos(y_proj)],
|
|
dim=1,
|
|
)
|
|
|
|
|
|
# =========================================================
|
|
# Regressor: Encoder + (xy [+ optional Fourier]) + FNO
|
|
# =========================================================
|
|
class FNOFieldRegressor_EncoderFNO(nn.Module):
|
|
def __init__(
|
|
self,
|
|
modes=6,
|
|
width=24,
|
|
blocks=4,
|
|
padding=0,
|
|
norm="gn",
|
|
dropout=0.10,
|
|
layerscale=0.10,
|
|
geom_ch=8,
|
|
use_fourier=False,
|
|
fourier_K=2,
|
|
fourier_scale=1.0,
|
|
dropout_xy=0.05,
|
|
):
|
|
super().__init__()
|
|
self.encoder = GeomEncoder(out_ch=geom_ch, drop=0.10)
|
|
|
|
self.use_fourier = bool(use_fourier)
|
|
self.fourier = FourierFeatures2D(K=fourier_K, scale=fourier_scale) if self.use_fourier else None
|
|
self.drop_xy = nn.Dropout2d(dropout_xy) if dropout_xy and dropout_xy > 0 else nn.Identity()
|
|
|
|
in_ch = geom_ch + 2 + (4 * fourier_K if self.use_fourier else 0)
|
|
|
|
self.fno = FNOModel2d(
|
|
modes=modes,
|
|
width=width,
|
|
blocks=blocks,
|
|
padding=padding,
|
|
in_channels=in_ch,
|
|
out_channels=2,
|
|
norm=norm,
|
|
dropout=dropout,
|
|
layerscale=layerscale,
|
|
)
|
|
|
|
self._cached_key = None
|
|
self._cached_xy = None # (1,2,Ny,Nx)
|
|
|
|
@staticmethod
|
|
def _sanitize_coords(coords: torch.Tensor) -> torch.Tensor:
|
|
# accept (B,M,2) or (B,1,M,2)
|
|
if coords.dim() == 4:
|
|
coords = coords[:, 0, :, :]
|
|
return coords
|
|
|
|
@staticmethod
|
|
def _coords_to_xy_grid(coords: torch.Tensor):
|
|
"""
|
|
coords: (B,M,2) or (M,2)
|
|
Build xy grid (1,2,Ny,Nx) using unique sorted x and y.
|
|
"""
|
|
coords0 = coords[0] if coords.dim() == 3 else coords # (M,2)
|
|
x = coords0[:, 0]
|
|
y = coords0[:, 1]
|
|
|
|
xs = torch.unique(x)
|
|
ys = torch.unique(y)
|
|
xs, _ = torch.sort(xs)
|
|
ys, _ = torch.sort(ys)
|
|
|
|
Ny, Nx = ys.numel(), xs.numel()
|
|
yy, xx = torch.meshgrid(ys, xs, indexing="ij")
|
|
xy = torch.stack([xx, yy], dim=0).unsqueeze(0) # (1,2,Ny,Nx)
|
|
return xy, int(Ny), int(Nx)
|
|
|
|
def forward(self, x_img: torch.Tensor, coords: torch.Tensor):
|
|
B = x_img.shape[0]
|
|
device = x_img.device
|
|
dtype = x_img.dtype
|
|
|
|
coords = self._sanitize_coords(coords).to(device=device)
|
|
xy, Ny, Nx = self._coords_to_xy_grid(coords)
|
|
|
|
key = (Ny, Nx, device.type, device.index, str(dtype))
|
|
if self._cached_key != key:
|
|
self._cached_key = key
|
|
self._cached_xy = xy.to(device=device, dtype=dtype)
|
|
|
|
xy_b = self._cached_xy.expand(B, -1, -1, -1)
|
|
xy_b = self.drop_xy(xy_b)
|
|
|
|
geom_feat = self.encoder(x_img, Ny, Nx)
|
|
|
|
if self.use_fourier:
|
|
ff = self.fourier(xy_b)
|
|
inp = torch.cat([geom_feat, xy_b, ff], dim=1)
|
|
else:
|
|
inp = torch.cat([geom_feat, xy_b], dim=1)
|
|
|
|
out = self.fno(inp) # (B,2,Ny,Nx)
|
|
out = out.permute(0, 2, 3, 1).contiguous().view(B, Ny * Nx, 2)
|
|
return out
|
|
|
|
|
|
# =========================================================
|
|
# Dataset loader: prefer build_dataset_full; else fallback concat
|
|
# =========================================================
|
|
def load_full_dataset_and_stats(matpath: str):
|
|
"""
|
|
For the new dataset we do NOT need norm_stats.
|
|
If build_dataset_full exists, it should return (full_ds, {}).
|
|
Otherwise we concat train/test from build_datasets.
|
|
"""
|
|
try:
|
|
from getdata import build_dataset_full # optional
|
|
full_ds, norm_stats = build_dataset_full(matpath)
|
|
if is_main_process():
|
|
print("[Dataset] Using build_dataset_full().")
|
|
return full_ds, norm_stats if norm_stats is not None else {}
|
|
except Exception as e:
|
|
if is_main_process():
|
|
print("[Dataset INFO] build_dataset_full() not found or failed. Using build_datasets() and concat.")
|
|
print(f" Reason: {repr(e)}")
|
|
train_ds, test_ds, _ = build_datasets(matpath)
|
|
full_ds = ConcatDataset([train_ds, test_ds])
|
|
return full_ds, {}
|
|
|
|
|
|
# =========================================================
|
|
# Indexed dataset wrapper (for per-sample saving)
|
|
# =========================================================
|
|
class IndexedDataset(torch.utils.data.Dataset):
|
|
"""
|
|
returns: (i, x_img, coords, E_true)
|
|
i is index within this ds (0..len(ds)-1)
|
|
"""
|
|
def __init__(self, ds):
|
|
self.ds = ds
|
|
|
|
def __len__(self):
|
|
return len(self.ds)
|
|
|
|
def __getitem__(self, i):
|
|
x_img, coords, E_true = self.ds[i]
|
|
return i, x_img, coords, E_true
|
|
|
|
|
|
# =========================================================
|
|
# Trainer
|
|
# =========================================================
|
|
class FNOOnly_DDP_Trainer:
|
|
def __init__(
|
|
self,
|
|
model,
|
|
device,
|
|
batch_size=32,
|
|
matpath="metalens_dataset.mat",
|
|
lr=2e-4,
|
|
weight_decay=3e-4,
|
|
num_workers=4,
|
|
pin_memory=True,
|
|
warmup_epochs=5,
|
|
min_lr_ratio=1e-3,
|
|
early_stop_patience=30,
|
|
grad_clip=1.0,
|
|
split_seed=20260303,
|
|
test_ratio=0.2,
|
|
):
|
|
self.device = device
|
|
self.matpath = matpath
|
|
self.warmup_epochs = int(warmup_epochs)
|
|
self.min_lr_ratio = float(min_lr_ratio)
|
|
self.early_stop_patience = int(early_stop_patience)
|
|
self.grad_clip = float(grad_clip)
|
|
self.split_seed = int(split_seed)
|
|
self.test_ratio = float(test_ratio)
|
|
|
|
self.model = model.to(self.device)
|
|
self.ddp_model = DDP(
|
|
self.model,
|
|
device_ids=[self.device.index],
|
|
output_device=self.device.index,
|
|
find_unused_parameters=False,
|
|
broadcast_buffers=False,
|
|
)
|
|
|
|
self.loss_fn = RelL2Loss()
|
|
self.optimizer = torch.optim.AdamW(self.ddp_model.parameters(), lr=lr, weight_decay=weight_decay)
|
|
|
|
self.scheduler = None
|
|
self.total_epochs = None
|
|
|
|
self.file_name = "fno_ddpsafe_final.pth"
|
|
self.best_name = "fno_ddpsafe_best.pth"
|
|
|
|
self.losses = []
|
|
|
|
self.train_set, self.test_set, self.norm_stats = self.load_dataset_ddpsafe()
|
|
|
|
self.train_sampler = DistributedSampler(self.train_set, shuffle=True, drop_last=False)
|
|
self.test_sampler = DistributedSampler(self.test_set, shuffle=False, drop_last=False)
|
|
|
|
self.batch_size = batch_size
|
|
self.num_workers = num_workers
|
|
self.pin_memory = pin_memory
|
|
|
|
self.train_loader = DataLoader(
|
|
self.train_set,
|
|
batch_size=batch_size,
|
|
sampler=self.train_sampler,
|
|
shuffle=False,
|
|
drop_last=False,
|
|
num_workers=num_workers,
|
|
pin_memory=pin_memory,
|
|
persistent_workers=(num_workers > 0),
|
|
)
|
|
|
|
self.test_loader = DataLoader(
|
|
self.test_set,
|
|
batch_size=batch_size,
|
|
sampler=self.test_sampler,
|
|
shuffle=False,
|
|
drop_last=False,
|
|
num_workers=num_workers,
|
|
pin_memory=pin_memory,
|
|
persistent_workers=(num_workers > 0),
|
|
)
|
|
|
|
if is_main_process():
|
|
Path("results").mkdir(parents=True, exist_ok=True)
|
|
|
|
def load_dataset_ddpsafe(self):
|
|
full_ds, norm_stats = load_full_dataset_and_stats(self.matpath)
|
|
N = len(full_ds)
|
|
|
|
n_test = int(round(N * self.test_ratio))
|
|
n_test = max(1, min(N - 1, n_test))
|
|
|
|
if is_main_process():
|
|
g = torch.Generator()
|
|
g.manual_seed(self.split_seed)
|
|
perm = torch.randperm(N, generator=g, dtype=torch.int64)
|
|
else:
|
|
perm = torch.empty(N, dtype=torch.int64)
|
|
|
|
perm = perm.to(self.device)
|
|
broadcast_tensor_(perm, src=0)
|
|
perm = perm.cpu()
|
|
|
|
test_idx = perm[:n_test].tolist()
|
|
train_idx = perm[n_test:].tolist()
|
|
|
|
train_ds = Subset(full_ds, train_idx)
|
|
test_ds = Subset(full_ds, test_idx)
|
|
|
|
if is_main_process():
|
|
print(
|
|
f"[DDP Split] Full={N}, Train={len(train_ds)}, Test={len(test_ds)}, "
|
|
f"seed={self.split_seed}, test_ratio={self.test_ratio}"
|
|
)
|
|
|
|
return train_ds, test_ds, norm_stats
|
|
|
|
def load_model(self, path):
|
|
try:
|
|
state = torch.load(path, map_location="cpu", weights_only=True)
|
|
except TypeError:
|
|
state = torch.load(path, map_location="cpu")
|
|
self.model.load_state_dict(state, strict=True)
|
|
|
|
@staticmethod
|
|
def _to_img_tensor(x_img: torch.Tensor) -> torch.Tensor:
|
|
# dataset gives (1,512,512); dataloader gives (B,1,512,512) already,
|
|
# but keep robust.
|
|
if x_img.dim() == 2:
|
|
x_img = x_img.unsqueeze(0).unsqueeze(0)
|
|
elif x_img.dim() == 3:
|
|
x_img = x_img.unsqueeze(0)
|
|
if x_img.size(1) != 1:
|
|
x_img = x_img[:, :1]
|
|
elif x_img.dim() == 4:
|
|
if x_img.size(1) != 1:
|
|
x_img = x_img[:, :1]
|
|
else:
|
|
raise ValueError(f"Unexpected x_img dim: {x_img.shape}")
|
|
return x_img
|
|
|
|
@staticmethod
|
|
def _sanitize_coords(coords: torch.Tensor) -> torch.Tensor:
|
|
# coords may become (B,M,2) from dataloader stacking; ok.
|
|
if coords.dim() == 4:
|
|
coords = coords[:, 0, :, :]
|
|
return coords
|
|
|
|
@staticmethod
|
|
def _sanitize_field(E_true: torch.Tensor) -> torch.Tensor:
|
|
# Want (B,M,2)
|
|
if E_true.dim() == 2:
|
|
E_true = E_true.unsqueeze(0)
|
|
if E_true.dim() == 4:
|
|
E_true = E_true[:, 0, :, :]
|
|
return E_true
|
|
|
|
def _build_scheduler(self):
|
|
assert self.total_epochs is not None
|
|
|
|
def lr_lambda(epoch: int):
|
|
if epoch < self.warmup_epochs:
|
|
return float(epoch + 1) / float(self.warmup_epochs)
|
|
t = (epoch - self.warmup_epochs) / max(1, (self.total_epochs - self.warmup_epochs))
|
|
return self.min_lr_ratio + (1.0 - self.min_lr_ratio) * 0.5 * (1.0 + np.cos(np.pi * t))
|
|
|
|
self.scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda=lr_lambda)
|
|
|
|
def get_data_loss(self, x_img, coords, E_true):
|
|
E_pred = self.ddp_model(x_img, coords)
|
|
return self.loss_fn(E_pred, E_true)
|
|
|
|
@torch.no_grad()
|
|
def test_loss(self):
|
|
self.ddp_model.eval()
|
|
loss_sum = torch.zeros((), device=self.device)
|
|
count = torch.zeros((), device=self.device)
|
|
|
|
for x_img, coords, E_true in self.test_loader:
|
|
x_img = self._to_img_tensor(x_img).to(self.device, non_blocking=True)
|
|
coords = self._sanitize_coords(coords).to(self.device, non_blocking=True)
|
|
E_true = self._sanitize_field(E_true).to(self.device, non_blocking=True)
|
|
|
|
loss = self.get_data_loss(x_img, coords, E_true)
|
|
bsz = x_img.shape[0]
|
|
loss_sum += loss.detach() * bsz
|
|
count += bsz
|
|
|
|
if world_size() > 1:
|
|
dist.all_reduce(loss_sum, op=dist.ReduceOp.SUM)
|
|
dist.all_reduce(count, op=dist.ReduceOp.SUM)
|
|
|
|
avg = (loss_sum / (count + 1e-12)).detach()
|
|
return float(avg.item())
|
|
|
|
def train(self, epochs=400):
|
|
self.total_epochs = int(epochs)
|
|
self._build_scheduler()
|
|
|
|
start_time = datetime.datetime.now()
|
|
best_test_loss = float("inf")
|
|
bad = 0
|
|
|
|
epoch_iter = range(self.total_epochs)
|
|
if is_main_process():
|
|
epoch_iter = tqdm(epoch_iter, desc="Training (DDP)")
|
|
|
|
for epoch in epoch_iter:
|
|
self.train_sampler.set_epoch(epoch)
|
|
self.ddp_model.train()
|
|
|
|
loss_sum = torch.zeros((), device=self.device)
|
|
count = torch.zeros((), device=self.device)
|
|
|
|
for x_img, coords, E_true in self.train_loader:
|
|
x_img = self._to_img_tensor(x_img).to(self.device, non_blocking=True)
|
|
coords = self._sanitize_coords(coords).to(self.device, non_blocking=True)
|
|
E_true = self._sanitize_field(E_true).to(self.device, non_blocking=True)
|
|
|
|
self.optimizer.zero_grad(set_to_none=True)
|
|
loss = self.get_data_loss(x_img, coords, E_true)
|
|
loss.backward()
|
|
|
|
if self.grad_clip and self.grad_clip > 0:
|
|
torch.nn.utils.clip_grad_norm_(self.ddp_model.parameters(), self.grad_clip)
|
|
|
|
self.optimizer.step()
|
|
|
|
bsz = x_img.shape[0]
|
|
loss_sum += loss.detach() * bsz
|
|
count += bsz
|
|
|
|
if world_size() > 1:
|
|
dist.all_reduce(loss_sum, op=dist.ReduceOp.SUM)
|
|
dist.all_reduce(count, op=dist.ReduceOp.SUM)
|
|
|
|
avg_train_loss = float((loss_sum / (count + 1e-12)).item())
|
|
avg_test_loss = self.test_loss()
|
|
self.scheduler.step()
|
|
|
|
stop_flag = torch.zeros((), device=self.device, dtype=torch.int32)
|
|
|
|
if is_main_process():
|
|
self.losses.append([epoch, avg_train_loss, avg_test_loss])
|
|
lr = self.optimizer.param_groups[0]["lr"]
|
|
|
|
improved = avg_test_loss < best_test_loss
|
|
if improved:
|
|
best_test_loss = avg_test_loss
|
|
bad = 0
|
|
torch.save(self.model.state_dict(), self.best_name)
|
|
else:
|
|
bad += 1
|
|
|
|
print(
|
|
f"Epoch {epoch} | train_loss={avg_train_loss:.6e} | "
|
|
f"test_loss={avg_test_loss:.6e} | best={best_test_loss:.6e} | "
|
|
f"bad={bad}/{self.early_stop_patience} | lr={lr:.3e}"
|
|
)
|
|
|
|
if bad >= self.early_stop_patience:
|
|
print(f"[Early stop] epoch={epoch}, best_test_loss={best_test_loss:.6e}")
|
|
stop_flag.fill_(1)
|
|
|
|
if world_size() > 1:
|
|
dist.broadcast(stop_flag, src=0)
|
|
|
|
if int(stop_flag.item()) == 1:
|
|
break
|
|
|
|
if world_size() > 1:
|
|
dist.barrier()
|
|
|
|
if is_main_process():
|
|
torch.save(self.model.state_dict(), self.file_name)
|
|
print("Best model saved to:", self.best_name)
|
|
print("Final model saved to:", self.file_name)
|
|
print("Training Time:", (datetime.datetime.now() - start_time).total_seconds(), "s")
|
|
|
|
# =========================================================
|
|
# Save one example (NO denorm for new dataset)
|
|
# =========================================================
|
|
def save_example_pred_true_mat_and_plot(
|
|
self,
|
|
idx: int = 0,
|
|
use_best: bool = True,
|
|
mat_save: str = "results/fno_ddpsafe_pred_true_example.mat",
|
|
fig_save: str = "results/fno_ddpsafe_field_plot_example.png",
|
|
interp_method: str = "cubic",
|
|
):
|
|
if not is_main_process():
|
|
return
|
|
|
|
if use_best and Path(self.best_name).exists():
|
|
self.load_model(self.best_name)
|
|
|
|
self.ddp_model.eval()
|
|
|
|
x_img, coords, E_true = self.test_set[idx]
|
|
|
|
if torch.is_tensor(coords) and coords.dim() == 3:
|
|
coords = coords[0, :]
|
|
coords_np = coords.detach().cpu().numpy()
|
|
|
|
if torch.is_tensor(E_true) and E_true.dim() == 3:
|
|
E_true = E_true[0, :, :]
|
|
E_true_np = E_true.detach().cpu().numpy()[None, ...] # (1,M,2)
|
|
|
|
with torch.no_grad():
|
|
x_img_t = self._to_img_tensor(x_img).to(self.device)
|
|
coords_t = self._sanitize_coords(coords.unsqueeze(0) if coords.dim() == 2 else coords).to(self.device)
|
|
E_pred_np = self.ddp_model(x_img_t, coords_t).detach().cpu().numpy() # (1,M,2)
|
|
|
|
Ere_pred, Eim_pred = E_pred_np[:, :, 0], E_pred_np[:, :, 1]
|
|
Ere_true, Eim_true = E_true_np[:, :, 0], E_true_np[:, :, 1]
|
|
|
|
savemat(
|
|
mat_save,
|
|
{
|
|
"Ere_pred": Ere_pred.astype(np.float32),
|
|
"Eim_pred": Eim_pred.astype(np.float32),
|
|
"Ere_true": Ere_true.astype(np.float32),
|
|
"Eim_true": Eim_true.astype(np.float32),
|
|
"coords_xy": coords_np.astype(np.float32),
|
|
},
|
|
)
|
|
print("Saved example mat:", mat_save)
|
|
|
|
# plot
|
|
X = coords_np[:, 0]
|
|
Y = coords_np[:, 1]
|
|
Xi, Yi = np.meshgrid(np.unique(X), np.unique(Y))
|
|
|
|
def interp(z):
|
|
Zi = griddata(points=(X, Y), values=z, xi=(Xi, Yi), method=interp_method)
|
|
from scipy.ndimage import gaussian_filter
|
|
return gaussian_filter(Zi, sigma=0.8)
|
|
|
|
re_err = np.abs(Ere_pred[0] - Ere_true[0])
|
|
im_err = np.abs(Eim_pred[0] - Eim_true[0])
|
|
|
|
fields = [
|
|
(interp(Ere_pred[0]), "E_real_pred"),
|
|
(interp(Ere_true[0]), "E_real_true"),
|
|
(interp(re_err), "E_real_error"),
|
|
(interp(Eim_pred[0]), "E_imag_pred"),
|
|
(interp(Eim_true[0]), "E_imag_true"),
|
|
(interp(im_err), "E_imag_error"),
|
|
]
|
|
|
|
fig, axes = plt.subplots(2, 3, figsize=(18, 9), sharex=True, sharey=True)
|
|
for ax, (Zi, title) in zip(axes.ravel(), fields):
|
|
pcm = ax.pcolormesh(Xi, Yi, Zi, shading="gouraud")
|
|
ax.set_title(title)
|
|
ax.set_xlabel("X")
|
|
ax.set_ylabel("Y")
|
|
ax.set_aspect("equal")
|
|
fig.colorbar(pcm, ax=ax, shrink=0.8)
|
|
|
|
plt.tight_layout()
|
|
plt.savefig(fig_save, dpi=200)
|
|
plt.close(fig)
|
|
print("Saved example plot:", fig_save)
|
|
|
|
# =========================================================
|
|
# Infer ALL train/test, save errors + ALL fields + coords
|
|
# =========================================================
|
|
@torch.no_grad()
|
|
def infer_all_and_save_errors_and_fields(
|
|
self,
|
|
use_best: bool = True,
|
|
save_path: str = "results/fno_ddpsafe_all_errors_and_fields.mat",
|
|
compute_mse_mae: bool = True,
|
|
save_coords_per_sample: bool = False,
|
|
dtype_save=np.float32,
|
|
):
|
|
if use_best and Path(self.best_name).exists():
|
|
if is_main_process():
|
|
print(f"[Infer] Loading best model: {self.best_name}")
|
|
self.load_model(self.best_name)
|
|
|
|
self.ddp_model.eval()
|
|
|
|
common_coords = None
|
|
if is_main_process():
|
|
_, coords0, _ = self.train_set[0]
|
|
if torch.is_tensor(coords0) and coords0.dim() == 3:
|
|
coords0 = coords0[0, :]
|
|
common_coords = coords0.detach().cpu().numpy().astype(dtype_save)
|
|
|
|
def build_infer_loader(ds):
|
|
ds_idx = IndexedDataset(ds)
|
|
sampler = DistributedSampler(ds_idx, shuffle=False, drop_last=False) # will pad -> we dedup by idx later
|
|
loader = DataLoader(
|
|
ds_idx,
|
|
batch_size=self.batch_size,
|
|
sampler=sampler,
|
|
shuffle=False,
|
|
drop_last=False,
|
|
num_workers=self.num_workers,
|
|
pin_memory=self.pin_memory,
|
|
persistent_workers=(self.num_workers > 0),
|
|
)
|
|
return loader, len(ds_idx)
|
|
|
|
def run_one_split(tag: str, ds):
|
|
loader, N = build_infer_loader(ds)
|
|
|
|
local_idx, local_rel = [], []
|
|
local_mse, local_mae = [], []
|
|
local_pred, local_true = [], []
|
|
local_coords = []
|
|
|
|
for idx, x_img, coords, E_true in loader:
|
|
x_img = self._to_img_tensor(x_img).to(self.device, non_blocking=True)
|
|
coords = self._sanitize_coords(coords).to(self.device, non_blocking=True)
|
|
E_true = self._sanitize_field(E_true).to(self.device, non_blocking=True)
|
|
|
|
E_pred = self.ddp_model(x_img, coords) # (B,M,2)
|
|
|
|
rel = rel_l2_per_sample(E_pred, E_true).detach().cpu().numpy()
|
|
|
|
if compute_mse_mae:
|
|
diff = (E_pred - E_true).detach()
|
|
mse = torch.mean(diff * diff, dim=(1, 2)).cpu().numpy()
|
|
mae = torch.mean(torch.abs(diff), dim=(1, 2)).cpu().numpy()
|
|
|
|
local_idx.append(idx.cpu().numpy().astype(np.int64))
|
|
local_rel.append(rel.astype(np.float64))
|
|
if compute_mse_mae:
|
|
local_mse.append(mse.astype(np.float64))
|
|
local_mae.append(mae.astype(np.float64))
|
|
|
|
local_pred.append(E_pred.detach().cpu().numpy().astype(dtype_save))
|
|
local_true.append(E_true.detach().cpu().numpy().astype(dtype_save))
|
|
|
|
if save_coords_per_sample:
|
|
local_coords.append(coords.detach().cpu().numpy().astype(dtype_save))
|
|
|
|
if len(local_idx) == 0:
|
|
pack = {
|
|
"idx": np.zeros((0,), np.int64),
|
|
"rel": np.zeros((0,), np.float64),
|
|
"mse": np.zeros((0,), np.float64) if compute_mse_mae else None,
|
|
"mae": np.zeros((0,), np.float64) if compute_mse_mae else None,
|
|
"pred": np.zeros((0, 0, 2), dtype_save),
|
|
"true": np.zeros((0, 0, 2), dtype_save),
|
|
"coords": np.zeros((0, 0, 2), dtype_save) if save_coords_per_sample else None,
|
|
"N": N,
|
|
}
|
|
else:
|
|
pack = {
|
|
"idx": np.concatenate(local_idx, axis=0),
|
|
"rel": np.concatenate(local_rel, axis=0),
|
|
"mse": np.concatenate(local_mse, axis=0) if compute_mse_mae else None,
|
|
"mae": np.concatenate(local_mae, axis=0) if compute_mse_mae else None,
|
|
"pred": np.concatenate(local_pred, axis=0),
|
|
"true": np.concatenate(local_true, axis=0),
|
|
"coords": np.concatenate(local_coords, axis=0) if save_coords_per_sample else None,
|
|
"N": N,
|
|
}
|
|
|
|
packs = all_gather_object(pack)
|
|
if not is_main_process():
|
|
return None
|
|
|
|
idx_all = np.concatenate([p["idx"] for p in packs], axis=0)
|
|
rel_all = np.concatenate([p["rel"] for p in packs], axis=0)
|
|
pred_all = np.concatenate([p["pred"] for p in packs], axis=0)
|
|
true_all = np.concatenate([p["true"] for p in packs], axis=0)
|
|
|
|
if compute_mse_mae:
|
|
mse_all = np.concatenate([p["mse"] for p in packs], axis=0)
|
|
mae_all = np.concatenate([p["mae"] for p in packs], axis=0)
|
|
else:
|
|
mse_all, mae_all = None, None
|
|
|
|
if save_coords_per_sample:
|
|
coords_all = np.concatenate([p["coords"] for p in packs if p["coords"] is not None], axis=0)
|
|
else:
|
|
coords_all = None
|
|
|
|
order = np.argsort(idx_all)
|
|
idx_all = idx_all[order]
|
|
rel_all = rel_all[order]
|
|
pred_all = pred_all[order]
|
|
true_all = true_all[order]
|
|
if compute_mse_mae:
|
|
mse_all = mse_all[order]
|
|
mae_all = mae_all[order]
|
|
if save_coords_per_sample:
|
|
coords_all = coords_all[order]
|
|
|
|
# dedup padded indices from DistributedSampler
|
|
_, first_pos = np.unique(idx_all, return_index=True)
|
|
idx_all = idx_all[first_pos]
|
|
rel_all = rel_all[first_pos]
|
|
pred_all = pred_all[first_pos]
|
|
true_all = true_all[first_pos]
|
|
if compute_mse_mae:
|
|
mse_all = mse_all[first_pos]
|
|
mae_all = mae_all[first_pos]
|
|
if save_coords_per_sample:
|
|
coords_all = coords_all[first_pos]
|
|
|
|
if idx_all.shape[0] != N:
|
|
missing = set(range(N)) - set(idx_all.tolist())
|
|
raise RuntimeError(
|
|
f"[{tag}] after dedup: got {idx_all.shape[0]} unique != N={N}, missing={len(missing)}"
|
|
)
|
|
|
|
out = {
|
|
f"{tag}_idx": idx_all.astype(np.int64),
|
|
f"{tag}_relL2": rel_all.astype(np.float64),
|
|
f"{tag}_E_pred": pred_all.astype(dtype_save),
|
|
f"{tag}_E_true": true_all.astype(dtype_save),
|
|
f"{tag}_relL2_mean": float(np.mean(rel_all)),
|
|
f"{tag}_relL2_std": float(np.std(rel_all)),
|
|
}
|
|
if compute_mse_mae:
|
|
out[f"{tag}_mse"] = mse_all.astype(np.float64)
|
|
out[f"{tag}_mae"] = mae_all.astype(np.float64)
|
|
out[f"{tag}_mse_mean"] = float(np.mean(mse_all))
|
|
out[f"{tag}_mae_mean"] = float(np.mean(mae_all))
|
|
if save_coords_per_sample:
|
|
out[f"{tag}_coords_xy_all"] = coords_all.astype(dtype_save)
|
|
|
|
return out
|
|
|
|
train_out = run_one_split("train", self.train_set)
|
|
test_out = run_one_split("test", self.test_set)
|
|
|
|
if is_main_process():
|
|
mdict = {}
|
|
mdict.update(train_out)
|
|
mdict.update(test_out)
|
|
|
|
if common_coords is not None:
|
|
mdict["coords_xy"] = common_coords # (M,2)
|
|
|
|
mdict["meta_matpath"] = self.matpath
|
|
mdict["meta_best_ckpt"] = self.best_name if use_best else ""
|
|
mdict["meta_world_size"] = int(world_size())
|
|
mdict["meta_compute_mse_mae"] = int(compute_mse_mae)
|
|
mdict["meta_save_coords_per_sample"] = int(save_coords_per_sample)
|
|
mdict["meta_dtype_save"] = str(dtype_save)
|
|
|
|
savemat(save_path, mdict)
|
|
print(f"[Infer] Saved ALL errors + ALL fields to: {save_path}")
|
|
print(f"[Infer] train_relL2_mean={mdict['train_relL2_mean']:.6e} | test_relL2_mean={mdict['test_relL2_mean']:.6e}")
|
|
|
|
|
|
# =========================================================
|
|
# Main
|
|
# =========================================================
|
|
def main():
|
|
device, local_rank = ddp_setup()
|
|
torch.backends.cudnn.benchmark = True
|
|
|
|
try:
|
|
# -------- config --------
|
|
epochs = 1000
|
|
batch_size = 32
|
|
|
|
# model
|
|
modes = 6
|
|
width = 24
|
|
blocks = 4
|
|
padding = 0
|
|
dropout = 0.10
|
|
layerscale = 0.10
|
|
geom_ch = 8
|
|
|
|
use_fourier = False
|
|
fourier_K = 2
|
|
fourier_scale = 1.0
|
|
dropout_xy = 0.05
|
|
|
|
# optim
|
|
lr = 2e-4
|
|
weight_decay = 3e-4
|
|
|
|
trainer = FNOOnly_DDP_Trainer(
|
|
model=FNOFieldRegressor_EncoderFNO(
|
|
modes=modes,
|
|
width=width,
|
|
blocks=blocks,
|
|
padding=padding,
|
|
norm="gn",
|
|
dropout=dropout,
|
|
layerscale=layerscale,
|
|
geom_ch=geom_ch,
|
|
use_fourier=use_fourier,
|
|
fourier_K=fourier_K,
|
|
fourier_scale=fourier_scale,
|
|
dropout_xy=dropout_xy,
|
|
),
|
|
device=device,
|
|
batch_size=batch_size,
|
|
matpath="metalens_dataset.mat", # ✅ NEW DATASET
|
|
lr=lr,
|
|
weight_decay=weight_decay,
|
|
num_workers=4,
|
|
pin_memory=True,
|
|
warmup_epochs=5,
|
|
min_lr_ratio=1e-3,
|
|
early_stop_patience=30,
|
|
grad_clip=1.0,
|
|
split_seed=20260303,
|
|
test_ratio=0.2,
|
|
)
|
|
|
|
trainer.train(epochs=epochs)
|
|
|
|
# ✅ 1) Save one example (NO denorm)
|
|
trainer.save_example_pred_true_mat_and_plot(
|
|
idx=0,
|
|
use_best=True,
|
|
mat_save="results/fno_ddpsafe_pred_true_example.mat",
|
|
fig_save="results/fno_ddpsafe_field_plot_example.png",
|
|
)
|
|
|
|
# ✅ 2) Save ALL train/test: errors + E_pred/E_true (+ coords)
|
|
trainer.infer_all_and_save_errors_and_fields(
|
|
use_best=True,
|
|
save_path="results/fno_ddpsafe_all_errors_and_fields.mat",
|
|
compute_mse_mae=True,
|
|
save_coords_per_sample=False, # coords are shared
|
|
dtype_save=np.float32,
|
|
)
|
|
|
|
finally:
|
|
ddp_cleanup()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |