File size: 4,707 Bytes
e9b87a5 | 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 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 | model:
start_epoch: 0
max_epoch: 10
lr: 5e-4
beta_1: 0.9
beta_2: 0.99
weight_decay: 1e-5
warmup_epochs: 10000
warmup_start_lr: 1e-8
eta_min: 1e-8
# ClimaX architecture params (from official config)
img_size: [32, 64] # 5.625° resolution (lat, lon)
patch_size: 2
embed_dim: 1024
depth: 8
decoder_depth: 2
num_heads: 16
mlp_ratio: 4.0
drop_path: 0.1
drop_rate: 0.1
# Forecasting settings
# NOTE: ERA5 data has 6-hour time steps (time_step=6 in HDF5 files).
# The lead_time passed to the model is (predict_range * hrs_each_step) / 100.
# For 6-hour forecast: predict_range=6, hrs_each_step=1 → lead_time=0.06
# For 72-hour forecast: predict_range=72, hrs_each_step=1 → lead_time=0.72
# The model conditions on lead_time; it must match the actual data gap.
predict_range: 6 # forecast lead time in hours
hrs_each_step: 1 # hours per data step factor for lead_time computation
checkpoint_dir: "./data/checkpoints"
patience: 50
# 整个数据读取流程
datapipe:
name: "ERA5"
task: "weather_forecasting"
# dataset设定
dataset:
type: "hdf5"
data_dir: './data/' # "$ONESCIENCE_DATASETS_DIR/ERA5/newh5/"
train_time: [2000, 2001]
val_time: [2002]
test_time: [2003]
img_size: [32, 64]
verbose: true
cache: false
# ClimaX 48 input variables: 3 constants + 3 surface + 42 pressure-level
channels:
# Static constants (3)
- "land_sea_mask"
- "orography"
- "lattitude"
# Surface variables (3)
- "2m_temperature"
- "10m_u_component_of_wind"
- "10m_v_component_of_wind"
# Geopotential at 7 pressure levels
- "geopotential_50"
- "geopotential_250"
- "geopotential_500"
- "geopotential_600"
- "geopotential_700"
- "geopotential_850"
- "geopotential_925"
# U component of wind at 7 pressure levels
- "u_component_of_wind_50"
- "u_component_of_wind_250"
- "u_component_of_wind_500"
- "u_component_of_wind_600"
- "u_component_of_wind_700"
- "u_component_of_wind_850"
- "u_component_of_wind_925"
# V component of wind at 7 pressure levels
- "v_component_of_wind_50"
- "v_component_of_wind_250"
- "v_component_of_wind_500"
- "v_component_of_wind_600"
- "v_component_of_wind_700"
- "v_component_of_wind_850"
- "v_component_of_wind_925"
# Temperature at 7 pressure levels
- "temperature_50"
- "temperature_250"
- "temperature_500"
- "temperature_600"
- "temperature_700"
- "temperature_850"
- "temperature_925"
# Relative humidity at 7 pressure levels
- "relative_humidity_50"
- "relative_humidity_250"
- "relative_humidity_500"
- "relative_humidity_600"
- "relative_humidity_700"
- "relative_humidity_850"
- "relative_humidity_925"
# Specific humidity at 7 pressure levels
- "specific_humidity_50"
- "specific_humidity_250"
- "specific_humidity_500"
- "specific_humidity_600"
- "specific_humidity_700"
- "specific_humidity_850"
- "specific_humidity_925"
# ClimaX output variables (5)
out_variables:
- "geopotential_500"
- "temperature_850"
- "2m_temperature"
- "10m_u_component_of_wind"
- "10m_v_component_of_wind"
# Short names for evaluation
variables:
- "z500"
- "t850"
- "t2m"
- "u10"
- "v10"
# 时间配置
time_range: ["2000-01-01", "2020-12-31"]
time_steps: 1
time_res: 6
# 空间配置
spatial_resolution: [5.625, 5.625]
# 采样配置
num_samples: -1 # -1 表示使用全部数据
shuffle: true
random_seed: 42
# 领域特定配置
extra:
levels: [500, 850]
lat_range: [-90, 90]
lon_range: [0, 360]
# 数据转换配置
transforms:
- type: "Normalize"
params:
mean: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
std: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
keys: ["input", "target"]
- type: "ToTensor"
params:
keys: null # null表示转换所有numpy数组
# DataLoader配置
dataloader:
mask_dtype: "float32"
batch_size: 1
num_workers: 1
pin_memory: true
drop_last: true
shuffle: false # 使用sampler时设为false
prefetch_factor: 2
persistent_workers: true
# 分布式配置
distributed:
enabled: true
sampler: "DistributedSampler"
rank: 0
world_size: 4
shuffle: true
seed: 42
drop_last: true
|