caowei
/
meta_cal
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
meta_cal/PI_code_test_FNO_inference.py

594 lines
20 KiB
Python
Raw Blame History

# 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()
Reference in New Issue View Git Blame Copy Permalink