scripts/inference.py

import sys
from pathlib import Path

# 获取项目根目录（inference.py上级的上级）
root_path = Path(__file__).parent.parent
sys.path.append(str(root_path))
import torch
import os
import glob
import numpy as np
import h5py
from tqdm import tqdm
from model.fgn import FGN
from onescience.utils.YParams import YParams
from onescience.datapipes.climate import ERA5Datapipe


def get_stats(data_dir, channels):
    """从新版 h5 中读取变量列表与归一化参数（均值/标准差）"""
    h5_files = sorted(glob.glob(os.path.join(data_dir, "data", "*.h5")))
    with h5py.File(h5_files[0], "r") as f:
        ds = f["fields"]
        all_variables = [v.decode() if isinstance(v, bytes) else v for v in ds.attrs["variables"]]
        mu = f["global_means"][:]   # [1, C, 1, 1]
        std = f["global_stds"][:]

    channel_indices = [all_variables.index(v) for v in channels]
    means = mu[:, channel_indices, :, :]
    stds = std[:, channel_indices, :, :]
    return means, stds


if __name__ == "__main__":
    current_path = os.getcwd()
    sys.path.append(current_path)

    ## Model config init
    config_file_path = os.path.join(current_path, "conf/config.yaml")
    cfg = YParams(config_file_path, "model")

    ## DataLoader init
    cfg_data = YParams(config_file_path, "datapipe")
    means, stds = get_stats(cfg_data.dataset.data_dir, cfg_data.dataset.channels)

    cfg['N_in_channels'] = len(cfg_data.dataset.channels)
    cfg['N_out_channels'] = len(cfg_data.dataset.channels)

    datapipe = ERA5Datapipe(
        dataset_dir=cfg_data.dataset.data_dir,
        used_variables=cfg_data.dataset.channels,
        used_years=cfg_data.dataset.test_time,
        distributed=False,
        input_steps=cfg.input_steps,
        output_steps=cfg.output_steps,
        batch_size=1,
        num_workers=4,
    )
    test_dataloader, _ = datapipe.get_dataloader("test")

    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    ckpt = torch.load(f"{cfg.checkpoint_dir}/model_bak.pth", map_location=device, weights_only=False)
    model = FGN(
        in_channels=cfg['N_in_channels'],
        out_channels=cfg['N_out_channels'],
        input_steps=cfg.input_steps,
        output_steps=cfg.output_steps,
        grid_shape=cfg.grid_shape,
        mesh_shape=cfg.mesh_shape,
        latent_dim=cfg.latent_dim,
        num_encoder_layers=cfg.num_encoder_layers,
        num_decoder_layers=cfg.num_decoder_layers,
        num_processor_blocks=cfg.num_processor_blocks,
        n_heads=cfg.n_heads,
        hidden_dim=cfg.hidden_dim,
        noise_dim=cfg.noise_dim,
        channel_weights=cfg.channel_weights,
    ).to(device)
    model.load_state_dict(ckpt["model_state_dict"])

    model.eval()
    os.makedirs('result/output/', exist_ok=True)
    print(f"📂 infer results will be generated to './result/output/'")
    print(f"📂 generating {cfg.num_members} ensemble members per init, saving their mean per frame")
    with torch.no_grad():
        for data in tqdm(test_dataloader, desc="Inferring testset", unit="batch"):
            invar = data[0].to(device, dtype=torch.float32)      # [1, input_steps, C, H, W]
            members = model(invar, num_members=cfg.num_members)  # [1, M, output_steps, C, H, W]
            pred = members[0].mean(dim=0).cpu().numpy()          # 集合成员均值 [T, C, H, W]
            for t in range(pred.shape[0]):
                fname = data[4][cfg.input_steps + t][0]          # 该预测帧对应的时刻
                pred_var = pred[t]                               # [C, H, W]
                pred_var = pred_var * stds + means
                np.save(f"result/output/{fname}.npy", pred_var)