File size: 4,127 Bytes
0ba2894 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | import sys
from pathlib import Path
# 获取项目根目录（train.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.prithvi_wxc import PrithviWxC
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=2,
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 = PrithviWxC(
in_channels=cfg['N_in_channels'],
input_size_time=cfg.input_size_time,
in_channels_static=cfg.in_channels_static,
n_lats_px=cfg.n_lats_px,
n_lons_px=cfg.n_lons_px,
patch_size_px=cfg.patch_size_px,
mask_unit_size_px=cfg.mask_unit_size_px,
mask_ratio_inputs=0.0,
embed_dim=cfg.embed_dim,
n_blocks_encoder=cfg.n_blocks_encoder,
n_blocks_decoder=cfg.n_blocks_decoder,
mlp_multiplier=cfg.mlp_multiplier,
n_heads=cfg.n_heads,
dropout=cfg.dropout,
drop_path=cfg.drop_path,
parameter_dropout=cfg.parameter_dropout,
residual=cfg.residual,
masking_mode=cfg.masking_mode,
positional_encoding=cfg.positional_encoding,
encoder_shifting=cfg.encoder_shifting,
decoder_shifting=cfg.decoder_shifting,
).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/'")
H, W = int(cfg.n_lats_px), int(cfg.n_lons_px)
static_path = os.path.join(cfg_data.dataset.data_dir, "static", "static.npy")
static_base = torch.from_numpy(np.load(static_path)).to(device=device, dtype=torch.float32).unsqueeze(0)
expected_static = (1, int(cfg.in_channels_static), H, W)
if tuple(static_base.shape) != expected_static:
raise ValueError(f"static data shape {tuple(static_base.shape)} != expected {expected_static}")
with torch.no_grad():
for data in tqdm(test_dataloader, desc="Inferring testset", unit="batch"):
invar = data[0].to(device, dtype=torch.float32) # [1, 2, C, H, W]
filename = data[4][-1][0]
B = invar.shape[0]
static = static_base.expand(B, -1, -1, -1)
lead_time = torch.full((B,), 6.0, device=device)
pred_var = model(invar, static, lead_time=lead_time).cpu().numpy()
pred_var = pred_var * stds + means
np.save(f"result/output/{filename}.npy", pred_var)
|