| # Prithvi WxC 训练配置示例 | |
| # 论文配置为 0.5°×0.625°（360/361×576）网格、160 动态变量、embed_dim=2560、 | |
| # 编码器 13 本地+12 全局块、解码器 3 本地+2 全局块（约 23 亿参数，需大规模显存）。 | |
| # 当前为连通性验证小配置（约百万级参数）。 | |
| model: | |
| start_epoch: 0 | |
| max_epoch: 100 | |
| lr: 1E-3 | |
| patience: 50 | |
| checkpoint_dir: "./data/checkpoints" | |
| # Prithvi WxC 结构参数 | |
| in_channels: 6 # 论文为 160（20 单层 + 10 变量 ×14 层） | |
| input_size_time: 2 # 两个输入时刻 | |
| in_channels_static: 4 # 静态通道数（论文 MEMRA-2 为 4） | |
| n_lats_px: 32 # 数据纬度方向像元数（论文约 360） | |
| n_lons_px: 64 # 数据经度方向像元数（论文 576） | |
| patch_size_px: [2, 2] # token 尺寸（论文 2×2 像素） | |
| mask_unit_size_px: [8, 8] # 掩码单元尺寸 | |
| mask_ratio_inputs: 0.0 # 预训练为 0.5；预报滚动微调为 0.0 | |
| embed_dim: 32 # 隐藏维度（论文 2560） | |
| n_blocks_encoder: 1 # 编码器本地-全局对数量（论文 13 对 -> 25 块） | |
| n_blocks_decoder: 1 # 解码器本地-全局对数量（论文 5 块） | |
| mlp_multiplier: 4.0 | |
| n_heads: 4 # 注意力头数（论文 16） | |
| dropout: 0.0 | |
| drop_path: 0.0 | |
| parameter_dropout: 0.0 | |
| residual: "none" # none / temporal / climate | |
| masking_mode: "global" | |
| positional_encoding: "absolute" | |
| encoder_shifting: false | |
| decoder_shifting: false | |
| # 整个数据读取流程 | |
| datapipe: | |
| name: "ERA5" | |
| task: "weather_forecasting" | |
| dataset: | |
| type: "hdf5" | |
| data_dir: './data/' | |
| train_time: [1951, 1952] | |
| val_time: [1953] | |
| test_time: [1954] | |
| img_size: [32, 64] | |
| verbose: true | |
| cache: false | |
| # 气象变量（论文 160 通道配置的子集，用于连通性验证） | |
| channels: ['10m_u_component_of_wind', '10m_v_component_of_wind', '2m_temperature', | |
| 'mean_sea_level_pressure', 'geopotential_500', 'temperature_850'] | |
| # DataLoader 配置 | |
| dataloader: | |
| mask_dtype: "float32" | |
| batch_size: 1 | |
| num_workers: 1 | |
| pin_memory: true | |
| drop_last: true | |
| shuffle: false | |
| prefetch_factor: 2 | |
| persistent_workers: true | |
| # 分布式配置 | |
| distributed: | |
| enabled: true | |
| sampler: "DistributedSampler" | |
| rank: 0 | |
| world_size: 2 | |
| shuffle: true | |
| seed: 42 | |
| drop_last: true |