From 29b4f7008d5b972052088395f0487a2d343ad010 Mon Sep 17 00:00:00 2001 From: caowei <2642367431@qq.com> Date: Tue, 23 Jun 2026 16:40:10 +0800 Subject: [PATCH] =?UTF-8?q?=E4=B8=8A=E4=BC=A0=E6=96=87=E4=BB=B6=E8=87=B3?= =?UTF-8?q?=20/?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- PI_code_test_FNO.py | 1041 +++++++++++++++++++++ PI_code_test_FNO_inference.py | 594 ++++++++++++ getdata.py | 154 +++ infer_ddp_save_all_dataset_with_inputs.py | 367 ++++++++ pido_test.sh | 22 + 5 files changed, 2178 insertions(+) create mode 100644 PI_code_test_FNO.py create mode 100644 PI_code_test_FNO_inference.py create mode 100644 getdata.py create mode 100644 infer_ddp_save_all_dataset_with_inputs.py create mode 100644 pido_test.sh diff --git a/PI_code_test_FNO.py b/PI_code_test_FNO.py new file mode 100644 index 0000000..0a3cb3a --- /dev/null +++ b/PI_code_test_FNO.py @@ -0,0 +1,1041 @@ +# 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() \ No newline at end of file diff --git a/PI_code_test_FNO_inference.py b/PI_code_test_FNO_inference.py new file mode 100644 index 0000000..190080f --- /dev/null +++ b/PI_code_test_FNO_inference.py @@ -0,0 +1,594 @@ +# PI_code_test_FNO_inference.py +# ========================================================= +# ✅ Inference ONLY for Encoder+FNO (no DDP, no training) +# ✅ Can read BOTH MATLAB v7 (scipy.io.loadmat) and v7.3 (HDF5 via h5py) +# ✅ Added: accurate inference timing (CUDA synchronized) + save timing into .mat +# +# Input .mat must contain (names can vary, see candidates below): +# - Rimg : (N,1,512,512) OR (512,512,N) OR (512,512) OR (1,512,512) +# - coords: (M,2) in [-1,1] (or (1,M,2)) +# +# Output .mat: +# - E_pred : (N,M,2) float32 (Re, Im) +# - coords_xy : (M,2) float32 +# - Rimg : (N,1,512,512) float32 +# - meta_* : timing + shapes + ckpt path +# ========================================================= + +import argparse +import math +import time +from pathlib import Path + +import numpy as np +import torch +import torch.nn as nn +from scipy.io import loadmat, savemat + + +# ------------------------- +# Utils +# ------------------------- +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) + + +def _pick_key(d, keys): + for k in keys: + if k in d: + return k + return None + + +def _as_numpy_from_h5(obj): + """ + Convert h5py dataset to numpy. + MATLAB v7.3 stores arrays in column-major and often with transposed dims; + We will reshape/transpose later in the Rimg handler. + """ + import numpy as _np + return _np.array(obj) + + +def load_mat_auto(mat_path: str): + """ + Load .mat file: + - MATLAB v7.2 and below: scipy.io.loadmat + - MATLAB v7.3 (HDF5): h5py + Returns: dict-like mapping name -> numpy array + """ + mat_path = str(mat_path) + try: + data = loadmat(mat_path) + # remove matlab meta keys if exist + return {k: v for k, v in data.items() if not k.startswith("__")} + except NotImplementedError: + import h5py + data = {} + with h5py.File(mat_path, "r") as f: + for k in f.keys(): + data[k] = _as_numpy_from_h5(f[k]) + return data + + +def load_meta_inference_mat(mat_in: str): + """ + Returns: + Rimg_np: (N,1,512,512) float32 + coords_np: (M,2) float32 + """ + data = load_mat_auto(mat_in) + + # candidates (you can add more if needed) + k_r = _pick_key(data, ["Rimg", "rimg", "geom", "geometry", "M"]) + k_c = _pick_key(data, ["coords", "coord", "coords_xy", "xy"]) + + if k_r is None or k_c is None: + raise KeyError( + f"[meta_inference] mat must contain geometry + coords.\n" + f"Found keys: {list(data.keys())}\n" + f"Geometry candidates: [Rimg,rimg,geom,geometry,M]\n" + f"Coords candidates: [coords,coord,coords_xy,xy]" + ) + + Rimg_np = np.array(data[k_r]) + coords_np = np.array(data[k_c]) + + # ---- coords -> (M,2) + coords_np = np.squeeze(coords_np) + + # MATLAB v7.3 (h5py) sometimes gives (2,M) for a (M,2) saved matrix + if coords_np.ndim == 2 and coords_np.shape[0] == 2 and coords_np.shape[1] != 2: + coords_np = coords_np.T + + if coords_np.ndim != 2 or coords_np.shape[1] != 2: + raise ValueError(f"[meta_inference] coords must be (M,2). Got {coords_np.shape}") + coords_np = coords_np.astype(np.float32) + + # ---- Rimg -> (N,1,512,512) + Rimg_np = np.squeeze(Rimg_np) + + if Rimg_np.ndim == 2: + # (512,512) + Rimg_np = Rimg_np[None, None, :, :] + + elif Rimg_np.ndim == 3: + # could be (512,512,N) OR (N,512,512) OR (1,512,512) + if Rimg_np.shape[0] == 1 and Rimg_np.shape[1] == 512 and Rimg_np.shape[2] == 512: + Rimg_np = Rimg_np[:, None, :, :] # (1,1,512,512) + + elif Rimg_np.shape[0] == 512 and Rimg_np.shape[1] == 512: + # (512,512,N) -> (N,1,512,512) + Rimg_np = np.transpose(Rimg_np, (2, 0, 1))[:, None, :, :] + + elif Rimg_np.shape[1] == 512 and Rimg_np.shape[2] == 512: + # (N,512,512) -> (N,1,512,512) + Rimg_np = Rimg_np[:, None, :, :] + + else: + raise ValueError(f"[meta_inference] Unrecognized Rimg 3D shape: {Rimg_np.shape}") + + elif Rimg_np.ndim == 4: + # could be (N,1,512,512) OR (1,512,512,N) OR (512,512,1,N) etc. + if Rimg_np.shape[1] == 1 and Rimg_np.shape[2] == 512 and Rimg_np.shape[3] == 512: + # (N,1,512,512) + pass + + elif Rimg_np.shape[0] == 1 and Rimg_np.shape[1] == 512 and Rimg_np.shape[2] == 512: + # (1,512,512,N) -> (N,1,512,512) + Rimg_np = np.transpose(Rimg_np, (3, 0, 1, 2)) + + elif Rimg_np.shape[0] == 512 and Rimg_np.shape[1] == 512 and Rimg_np.shape[2] == 1: + # (512,512,1,N) -> (N,1,512,512) + Rimg_np = np.transpose(Rimg_np, (3, 2, 0, 1)) + + else: + raise ValueError(f"[meta_inference] Unrecognized Rimg 4D shape: {Rimg_np.shape}") + + else: + raise ValueError(f"[meta_inference] Unrecognized Rimg shape: {Rimg_np.shape}") + + if Rimg_np.shape[1] != 1 or Rimg_np.shape[-2:] != (512, 512): + raise ValueError(f"[meta_inference] Rimg must be (N,1,512,512). Got {Rimg_np.shape}") + + return Rimg_np.astype(np.float32), coords_np.astype(np.float32) + + +# ------------------------- +# FNO core +# ------------------------- +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 + + +# ------------------------- +# 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 (OFF by default) +# ------------------------- +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: + if coords.dim() == 4: + coords = coords[:, 0, :, :] + return coords + + @staticmethod + def _coords_to_xy_grid(coords: torch.Tensor): + 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, getattr(device, "index", None), 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 + + +# ------------------------- +# Inference +# ------------------------- +@torch.no_grad() +def run_inference( + mat_in: str, + ckpt: str, + mat_out: str, + device: str, + batch_size: int, + modes: int, + width: int, + blocks: int, + padding: int, + dropout: float, + layerscale: float, + geom_ch: int, + use_fourier: bool, + fourier_K: int, + fourier_scale: float, + dropout_xy: float, +): + Rimg_np, coords_np = load_meta_inference_mat(mat_in) + N = int(Rimg_np.shape[0]) + M = int(coords_np.shape[0]) + + dev = torch.device(device if torch.cuda.is_available() else "cpu") + + 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, + ).to(dev) + + # load weights + try: + state = torch.load(ckpt, map_location="cpu", weights_only=True) + except TypeError: + state = torch.load(ckpt, map_location="cpu") + model.load_state_dict(state, strict=True) + model.eval() + + coords_t = torch.from_numpy(coords_np).to(dev).unsqueeze(0) # (1,M,2) + + # ---- timing helpers (CUDA safe) + def _sync(): + if dev.type == "cuda": + torch.cuda.synchronize(dev) + + # warmup (recommended) + warmup_steps = min(2, (N + batch_size - 1) // batch_size) + for wi in range(warmup_steps): + s = wi * batch_size + e = min(N, s + batch_size) + x_img = torch.from_numpy(Rimg_np[s:e]).to(dev, dtype=torch.float32) + coords_b = coords_t.expand(e - s, -1, -1).contiguous() + _ = model(x_img, coords_b) + _sync() + + preds = [] + t_total0 = time.perf_counter() + forward_time_sum = 0.0 + num_samples_done = 0 + + for s in range(0, N, batch_size): + e = min(N, s + batch_size) + x_img = torch.from_numpy(Rimg_np[s:e]).to(dev, dtype=torch.float32) + coords_b = coords_t.expand(e - s, -1, -1).contiguous() + + _sync() + t0 = time.perf_counter() + pred = model(x_img, coords_b) # (B,M,2) + _sync() + t1 = time.perf_counter() + + forward_time_sum += (t1 - t0) + num_samples_done += (e - s) + + preds.append(pred.detach().cpu().numpy().astype(np.float32)) + + t_total1 = time.perf_counter() + total_time = t_total1 - t_total0 + + avg_time_per_sample = total_time / max(num_samples_done, 1) + avg_forward_per_sample = forward_time_sum / max(num_samples_done, 1) + + fps_total = num_samples_done / max(total_time, 1e-12) + fps_forward = num_samples_done / max(forward_time_sum, 1e-12) + + E_pred = np.concatenate(preds, axis=0) # (N,M,2) + + print("========== Inference Timing ==========") + print(f"[Time] total elapsed : {total_time:.6f} s") + print(f"[Time] total forward (model) : {forward_time_sum:.6f} s") + print(f"[Time] avg elapsed / sample : {avg_time_per_sample*1e3:.3f} ms") + print(f"[Time] avg forward / sample : {avg_forward_per_sample*1e3:.3f} ms") + print(f"[Perf] throughput (total) : {fps_total:.3f} samples/s") + print(f"[Perf] throughput (forward) : {fps_forward:.3f} samples/s") + print("======================================") + + savemat( + mat_out, + { + "E_pred": E_pred, + "coords_xy": coords_np.astype(np.float32), + "meta_ckpt": str(ckpt), + "meta_N": np.int64(N), + "meta_M": np.int64(M), + + # timing meta + "meta_time_total_s": np.float64(total_time), + "meta_time_forward_s": np.float64(forward_time_sum), + "meta_ms_per_sample_total": np.float64(avg_time_per_sample * 1e3), + "meta_ms_per_sample_forward": np.float64(avg_forward_per_sample * 1e3), + "meta_throughput_total_sps": np.float64(fps_total), + "meta_throughput_forward_sps": np.float64(fps_forward), + }, + ) + + print(f"[OK] mat_in : {mat_in}") + print(f"[OK] ckpt : {ckpt}") + print(f"[OK] mat_out : {mat_out}") + print(f"[OK] shapes : Rimg={Rimg_np.shape}, coords={coords_np.shape}, E_pred={E_pred.shape}") + + +def parse_args(): + p = argparse.ArgumentParser("Inference only: Encoder+FNO") + p.add_argument("--mat_in", type=str, default="meta_inference.mat") + p.add_argument("--ckpt", type=str, default="fno_ddpsafe_best.pth") + p.add_argument("--mat_out", type=str, default="meta_inference_pred.mat") + p.add_argument("--device", type=str, default="cuda:0") + p.add_argument("--batch_size", type=int, default=16) + + # must match training hyperparameters + p.add_argument("--modes", type=int, default=6) + p.add_argument("--width", type=int, default=24) + p.add_argument("--blocks", type=int, default=4) + p.add_argument("--padding", type=int, default=0) + p.add_argument("--dropout", type=float, default=0.10) + p.add_argument("--layerscale", type=float, default=0.10) + p.add_argument("--geom_ch", type=int, default=8) + + p.add_argument("--use_fourier", action="store_true") + p.add_argument("--fourier_K", type=int, default=2) + p.add_argument("--fourier_scale", type=float, default=1.0) + p.add_argument("--dropout_xy", type=float, default=0.05) + return p.parse_args() + + +def main(): + args = parse_args() + + if not Path(args.mat_in).exists(): + raise FileNotFoundError(f"Input mat not found: {args.mat_in}") + if not Path(args.ckpt).exists(): + raise FileNotFoundError(f"Checkpoint not found: {args.ckpt}") + + run_inference( + mat_in=args.mat_in, + ckpt=args.ckpt, + mat_out=args.mat_out, + device=args.device, + batch_size=args.batch_size, + modes=args.modes, + width=args.width, + blocks=args.blocks, + padding=args.padding, + dropout=args.dropout, + layerscale=args.layerscale, + geom_ch=args.geom_ch, + use_fourier=bool(args.use_fourier), + fourier_K=args.fourier_K, + fourier_scale=args.fourier_scale, + dropout_xy=args.dropout_xy, + ) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/getdata.py b/getdata.py new file mode 100644 index 0000000..85b8b81 --- /dev/null +++ b/getdata.py @@ -0,0 +1,154 @@ +import h5py +import numpy as np +import torch +from torch.utils.data import Dataset + + +# ---------------------------- +# Helpers: robust mat(v7.3) read +# ---------------------------- +def _read_array(f, key): + if key not in f: + raise KeyError(f"Key '{key}' not found. Available keys: {list(f.keys())}") + return np.array(f[key]) + + +def _ensure_coords_m2(coords): + # coords expected: (m,2). Some mat files read as (2,m) + if coords.ndim != 2: + raise ValueError(f"coords must be 2D, got {coords.shape}") + if coords.shape[1] == 2: + return coords + if coords.shape[0] == 2: + return coords.T + raise ValueError(f"coords should be (m,2) or (2,m), got {coords.shape}") + + +def _ensure_X_nchw(X, H=512, W=512): + """ + Expect geometry: + (N,1,H,W) OR (H,W,1,N) OR (1,H,W,N) + """ + if X.ndim != 4: + raise ValueError(f"X must be 4D, got {X.shape}") + + # (N,1,H,W) + if X.shape[0] > 1 and X.shape[1] == 1 and X.shape[2] == H and X.shape[3] == W: + return X + + # (H,W,1,N) -> (N,1,H,W) + if X.shape[0] == H and X.shape[1] == W and X.shape[2] == 1: + return np.transpose(X, (3, 2, 0, 1)) + + # (1,H,W,N) -> (N,1,H,W) + if X.shape[0] == 1 and X.shape[1] == H and X.shape[2] == W: + return np.transpose(X, (3, 0, 1, 2)) + + raise ValueError(f"Unrecognized X layout: {X.shape}") + + +def _ensure_Y_Nm(Y, N_expected=None, m_expected=None): + # expect (N,m) or (m,N) + if Y.ndim != 2: + raise ValueError(f"Y must be 2D, got {Y.shape}") + + if N_expected is not None and Y.shape[0] == N_expected: + return Y + if N_expected is not None and Y.shape[1] == N_expected: + return Y.T + + if m_expected is not None and Y.shape[1] == m_expected: + return Y + if m_expected is not None and Y.shape[0] == m_expected: + return Y.T + + return Y + + +# ---------------------------- +# Dataset (NO normalization here!) +# ---------------------------- +class MetalensDataset(Dataset): + def __init__(self, X_geom, coords, Y_re, Y_im): + """ + X_geom: (N,1,512,512) already in {0,1} or float + coords: (m,2) already normalized to [-1,1] + Y_re: (N,m) already normalized + Y_im: (N,m) already normalized + """ + # cast geometry to float32 for convs + self.X = torch.tensor(X_geom, dtype=torch.float32) + self.coords = torch.tensor(coords, dtype=torch.float32) # shared for all samples + self.Yre = torch.tensor(Y_re, dtype=torch.float32) + self.Yim = torch.tensor(Y_im, dtype=torch.float32) + + def __len__(self): + return self.X.shape[0] + + def __getitem__(self, idx): + x = self.X[idx] # (1,512,512) + y = torch.stack([self.Yre[idx], self.Yim[idx]], dim=-1) # (m,2) + return x, self.coords, y + + +# ---------------------------- +# Build datasets +# ---------------------------- +def build_datasets( + mat_path="metalens_dataset.mat", + train_ratio=0.8, + keys=None, +): + """ + Read your new dataset: + Rimg : (N,1,512,512) (or equivalent layout) + coords : (m,2) + Ere_flat : (N,m) + Eim_flat : (N,m) + No norm_stats needed. + """ + if keys is None: + keys = dict( + X="Rimg", + coords="coords", + Yre="Ere_flat", + Yim="Eim_flat", + ) + + with h5py.File(mat_path, "r") as f: + X = _read_array(f, keys["X"]) + coords = _read_array(f, keys["coords"]) + Yre = _read_array(f, keys["Yre"]) + Yim = _read_array(f, keys["Yim"]) + + coords = _ensure_coords_m2(coords) + X = _ensure_X_nchw(X, 512, 512) + + N = X.shape[0] + m = coords.shape[0] + + Yre = _ensure_Y_Nm(Yre, N_expected=N, m_expected=m) + Yim = _ensure_Y_Nm(Yim, N_expected=N, m_expected=m) + + if Yre.shape != (N, m) or Yim.shape != (N, m): + raise ValueError(f"Y shapes mismatch: Yre={Yre.shape}, Yim={Yim.shape}, expected {(N, m)}") + + if N < 2: + raise ValueError(f"N too small: {N}") + + n_train = int(train_ratio * N) + n_train = min(max(1, n_train), N - 1) + + X_train, X_test = X[:n_train], X[n_train:] + Yre_train, Yre_test = Yre[:n_train], Yre[n_train:] + Yim_train, Yim_test = Yim[:n_train], Yim[n_train:] + + train_ds = MetalensDataset(X_train, coords, Yre_train, Yim_train) + test_ds = MetalensDataset(X_test, coords, Yre_test, Yim_test) + + print(f"[build_datasets] X: {X.shape}, coords: {coords.shape}, Y: {Yre.shape}") + print(f"[build_datasets] Split train={len(train_ds)}, test={len(test_ds)}") + + # norm_stats 不再需要：返回空 dict + norm_stats = {} + return train_ds, test_ds, norm_stats \ No newline at end of file diff --git a/infer_ddp_save_all_dataset_with_inputs.py b/infer_ddp_save_all_dataset_with_inputs.py new file mode 100644 index 0000000..35cb1e5 --- /dev/null +++ b/infer_ddp_save_all_dataset_with_inputs.py @@ -0,0 +1,367 @@ +# infer_ddp_save_all_dataset_with_inputs.py +# ========================================================= +# ✅ DDP / single-GPU inference for FULL dataset +# ✅ Save per-sample: global_idx, Rimg(input), coords, E_pred, E_true, errors +# ✅ DDP-safe: each rank writes shard -> rank0 merges, sorts by idx, dedups +# Output: results/full_infer_with_inputs.h5 +# ========================================================= + +import os +import shutil +import numpy as np +import torch +import torch.distributed as dist +from torch.utils.data import DataLoader, ConcatDataset +from torch.utils.data.distributed import DistributedSampler +from pathlib import Path +import h5py + +# --------------------------------------------------------- +# Import from your training script (must be in same folder or PYTHONPATH) +# It must have "if __name__ == '__main__': main()" guard (you already do). +# --------------------------------------------------------- +from PI_code_test_FNO import ( + ddp_setup, + ddp_cleanup, + is_main_process, + world_size, + broadcast_tensor_, + load_full_dataset_and_stats, + IndexedDataset, + rel_l2_per_sample, + FNOFieldRegressor_EncoderFNO, +) + + +# --------------------------------------------------------- +# Small helpers (same behavior as trainer) +# --------------------------------------------------------- +def _to_img_tensor(x_img: torch.Tensor) -> torch.Tensor: + # want (B,1,512,512) + 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: {tuple(x_img.shape)}") + return x_img + + +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 + + +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 load_model_strict(model: torch.nn.Module, ckpt_path: str): + ckpt_path = str(ckpt_path) + try: + state = torch.load(ckpt_path, map_location="cpu", weights_only=True) + except TypeError: + state = torch.load(ckpt_path, map_location="cpu") + model.load_state_dict(state, strict=True) + + +# --------------------------------------------------------- +# Main inference +# --------------------------------------------------------- +@torch.no_grad() +def main(): + device, local_rank = ddp_setup() + torch.backends.cudnn.benchmark = True + + try: + # ========================= + # Config (edit as needed) + # ========================= + matpath = "metalens_dataset.mat" + + # ckpt + ckpt_path = "fno_ddpsafe_best.pth" # or "fno_ddpsafe_final.pth" + + # dataloader + batch_size = 16 # inference: you can increase if memory allows + num_workers = 4 + pin_memory = True + + # output + out_dir = Path("results") + out_dir.mkdir(parents=True, exist_ok=True) + shards_dir = out_dir / "infer_shards" + final_h5 = out_dir / "full_infer_with_inputs.h5" + + # ========================= + # Build FULL dataset + # ========================= + full_ds, _ = load_full_dataset_and_stats(matpath) + N = len(full_ds) + + if is_main_process(): + print(f"[Infer] FULL dataset size: N={N}") + + # ========================= + # Build model (must match training) + # ========================= + model = FNOFieldRegressor_EncoderFNO( + 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, + ).to(device) + + # DDP wrap + ddp_model = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[device.index], + output_device=device.index, + find_unused_parameters=False, + broadcast_buffers=False, + ) + + if is_main_process(): + print(f"[Infer] Loading ckpt: {ckpt_path}") + load_model_strict(model, ckpt_path) + ddp_model.eval() + + # ========================= + # coords (shared) save once (rank0) + # ========================= + coords_shared = None + if is_main_process(): + x0, coords0, E0 = full_ds[0] + if torch.is_tensor(coords0) and coords0.dim() == 3: + coords0 = coords0[0, :] + coords_shared = coords0.detach().cpu().numpy().astype(np.float32) + + # ========================= + # DataLoader with indices + # ========================= + ds_idx = IndexedDataset(full_ds) # returns: (i, x_img, coords, E_true) + sampler = DistributedSampler(ds_idx, shuffle=False, drop_last=False) + loader = DataLoader( + ds_idx, + batch_size=batch_size, + sampler=sampler, + shuffle=False, + drop_last=False, + num_workers=num_workers, + pin_memory=pin_memory, + persistent_workers=(num_workers > 0), + ) + + # ========================= + # Each rank writes its shard to avoid giant all_gather + # ========================= + if is_main_process(): + if shards_dir.exists(): + shutil.rmtree(shards_dir) + shards_dir.mkdir(parents=True, exist_ok=True) + if world_size() > 1: + dist.barrier() + + shard_path = shards_dir / f"shard_rank{dist.get_rank() if dist.is_initialized() else 0}.h5" + if shard_path.exists(): + shard_path.unlink() + + # We append to HDF5 with resizable datasets + with h5py.File(shard_path, "w") as f: + # Create empty resizable datasets + f.create_dataset("idx", shape=(0,), maxshape=(None,), dtype="int64", chunks=True) + f.create_dataset("Rimg", shape=(0, 1, 512, 512), maxshape=(None, 1, 512, 512), + dtype="float32", chunks=(1, 1, 512, 512)) + # E_pred/E_true are (B,M,2) => store as (n, M, 2) + # We don't know M at creation time until first batch -> create lazily + dset_Ep = None + dset_Et = None + f.create_dataset("relL2", shape=(0,), maxshape=(None,), dtype="float64", chunks=True) + f.create_dataset("mse", shape=(0,), maxshape=(None,), dtype="float64", chunks=True) + f.create_dataset("mae", shape=(0,), maxshape=(None,), dtype="float64", chunks=True) + + n_written = 0 + for idx, x_img, coords, E_true in loader: + idx_np = idx.cpu().numpy().astype(np.int64) + + x_img = _to_img_tensor(x_img).to(device, non_blocking=True) + coords = _sanitize_coords(coords).to(device, non_blocking=True) + E_true = _sanitize_field(E_true).to(device, non_blocking=True) + + E_pred = ddp_model(x_img, coords) # (B,M,2) + + rel = rel_l2_per_sample(E_pred, E_true).detach().cpu().numpy().astype(np.float64) + diff = (E_pred - E_true).detach() + mse = torch.mean(diff * diff, dim=(1, 2)).cpu().numpy().astype(np.float64) + mae = torch.mean(torch.abs(diff), dim=(1, 2)).cpu().numpy().astype(np.float64) + + x_np = x_img.detach().cpu().numpy().astype(np.float32) # (B,1,512,512) + Ep_np = E_pred.detach().cpu().numpy().astype(np.float32) # (B,M,2) + Et_np = E_true.detach().cpu().numpy().astype(np.float32) # (B,M,2) + + B = idx_np.shape[0] + M = Ep_np.shape[1] + + # Lazy create E datasets once we know M + if dset_Ep is None: + dset_Ep = f.create_dataset( + "E_pred", shape=(0, M, 2), maxshape=(None, M, 2), + dtype="float32", chunks=(1, M, 2) + ) + dset_Et = f.create_dataset( + "E_true", shape=(0, M, 2), maxshape=(None, M, 2), + dtype="float32", chunks=(1, M, 2) + ) + + # Resize and append + new_n = n_written + B + f["idx"].resize((new_n,)) + f["Rimg"].resize((new_n, 1, 512, 512)) + f["relL2"].resize((new_n,)) + f["mse"].resize((new_n,)) + f["mae"].resize((new_n,)) + dset_Ep.resize((new_n, M, 2)) + dset_Et.resize((new_n, M, 2)) + + f["idx"][n_written:new_n] = idx_np + f["Rimg"][n_written:new_n, ...] = x_np + f["relL2"][n_written:new_n] = rel + f["mse"][n_written:new_n] = mse + f["mae"][n_written:new_n] = mae + dset_Ep[n_written:new_n, ...] = Ep_np + dset_Et[n_written:new_n, ...] = Et_np + + n_written = new_n + + if world_size() > 1: + dist.barrier() + + # ========================= + # Rank0 merges shards -> final_h5 + # ========================= + if is_main_process(): + print(f"[Infer] Shards saved in: {shards_dir}") + + shard_files = sorted(shards_dir.glob("shard_rank*.h5")) + if len(shard_files) == 0: + raise RuntimeError("No shard files found.") + + # Load all shard arrays (still much safer than all_gather on GPU) + all_idx = [] + all_R = [] + all_rel = [] + all_mse = [] + all_mae = [] + all_Ep = [] + all_Et = [] + + M_final = None + for sf in shard_files: + with h5py.File(sf, "r") as f: + idx = f["idx"][:] + Rimg = f["Rimg"][:] + rel = f["relL2"][:] + mse = f["mse"][:] + mae = f["mae"][:] + Ep = f["E_pred"][:] + Et = f["E_true"][:] + + if M_final is None: + M_final = Ep.shape[1] + else: + if Ep.shape[1] != M_final: + raise RuntimeError(f"M mismatch among shards: {sf} has M={Ep.shape[1]} vs {M_final}") + + all_idx.append(idx) + all_R.append(Rimg) + all_rel.append(rel) + all_mse.append(mse) + all_mae.append(mae) + all_Ep.append(Ep) + all_Et.append(Et) + + idx_all = np.concatenate(all_idx, axis=0) + R_all = np.concatenate(all_R, axis=0) + rel_all = np.concatenate(all_rel, axis=0) + mse_all = np.concatenate(all_mse, axis=0) + mae_all = np.concatenate(all_mae, axis=0) + Ep_all = np.concatenate(all_Ep, axis=0) + Et_all = np.concatenate(all_Et, axis=0) + + # Sort by idx + order = np.argsort(idx_all) + idx_all = idx_all[order] + R_all = R_all[order] + rel_all = rel_all[order] + mse_all = mse_all[order] + mae_all = mae_all[order] + Ep_all = Ep_all[order] + Et_all = Et_all[order] + + # Dedup (DistributedSampler may pad) + uniq_idx, first_pos = np.unique(idx_all, return_index=True) + idx_all = idx_all[first_pos] + R_all = R_all[first_pos] + rel_all = rel_all[first_pos] + mse_all = mse_all[first_pos] + mae_all = mae_all[first_pos] + Ep_all = Ep_all[first_pos] + Et_all = Et_all[first_pos] + + if idx_all.shape[0] != N: + missing = set(range(N)) - set(idx_all.tolist()) + raise RuntimeError(f"[Merge] unique={idx_all.shape[0]} != N={N}, missing={len(missing)}") + + # Save final H5 + if final_h5.exists(): + final_h5.unlink() + + with h5py.File(final_h5, "w") as f: + f.create_dataset("idx", data=idx_all.astype(np.int64)) + f.create_dataset("Rimg", data=R_all.astype(np.float32), compression="gzip", compression_opts=4) + f.create_dataset("coords_xy", data=coords_shared.astype(np.float32)) + f.create_dataset("E_pred", data=Ep_all.astype(np.float32), compression="gzip", compression_opts=4) + f.create_dataset("E_true", data=Et_all.astype(np.float32), compression="gzip", compression_opts=4) + f.create_dataset("relL2", data=rel_all.astype(np.float64)) + f.create_dataset("mse", data=mse_all.astype(np.float64)) + f.create_dataset("mae", data=mae_all.astype(np.float64)) + + f.attrs["meta_matpath"] = str(matpath) + f.attrs["meta_ckpt"] = str(ckpt_path) + f.attrs["meta_world_size"] = int(world_size()) + f.attrs["N"] = int(N) + f.attrs["M"] = int(Ep_all.shape[1]) + + print(f"[Infer] Final saved: {final_h5}") + print(f"[Infer] relL2 mean={rel_all.mean():.6e}, std={rel_all.std():.6e}") + print(f"[Infer] mse mean={mse_all.mean():.6e} | mae mean={mae_all.mean():.6e}") + + # Optional: remove shards + # shutil.rmtree(shards_dir) + + finally: + ddp_cleanup() + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/pido_test.sh b/pido_test.sh new file mode 100644 index 0000000..1d7c891 --- /dev/null +++ b/pido_test.sh @@ -0,0 +1,22 @@ +#!/bin/bash +#SBATCH --job-name=run_node05 # 作业名称 +#SBATCH --output=pidn_output_%j # 标准输出文件（%j 会被替换为作业 ID） +#SBATCH --error=pidn_error_%j.txt # 标准错误文件（%j 会被替换为作业 ID） +#SBATCH --time=100:00:00 # 运行时间限制 +#SBATCH --partition=gpu # 请求 GPU 分区 +#SBATCH --cpus-per-task=10 # 节点请求的cpu核心数 +#SBATCH --mem=128G # 请求内存大小 +#SBATCH --gres=gpu:7 # 请求 4 个 GPU 资源 +#SBATCH --nodelist=node06 # 指定使用 node06 节点 + + +# 激活虚拟环境 +eval "$(/public/apps/miniconda3/bin/conda shell.bash hook)" +conda activate deepnet + +# 切换到存放 Python 脚本的目录 +cd /public/home/cw/meta_cal/20260304_FNO_512 + + +# 执行 Python 脚本 +torchrun --nproc_per_node=7 PI_code_test_FNO.py \ No newline at end of file