367 lines
14 KiB
Python
367 lines
14 KiB
Python
# 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() |