上传文件至 /

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()

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

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()

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