File size: 5,265 Bytes
c649472
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
work_dir: ./results
ckpt_dir: ./results/checkpoints/
log_dir: ./results/logs/
print_config: true
ignore_warnings: true
test_after_training: false
save_config_to_wandb: true
verbose: true
seed: 11
torch_matmul_precision: high
name: FV3GFS-Ipol6h
name_suffix: null
suffix: null
trainer:
  profiler: {}
  _target_: pytorch_lightning.Trainer
  accelerator: gpu
  devices: 4
  strategy: ddp_find_unused_parameters_false
  min_epochs: 1
  max_epochs: 60
  gradient_clip_val: 0.5
  accumulate_grad_batches: 3
  precision: 16
  num_sanity_val_steps: 1
  val_check_interval: 1.0
  deterministic: true
  log_every_n_steps: 50
model:
  _target_: src.models.sfno.sfnonet.SphericalFourierNeuralOperatorNet
  embed_dim: 256
  spectral_transform: sht
  filter_type: linear
  operator_type: dhconv
  num_layers: 8
  use_mlp: true
  mlp_ratio: 2.0
  scale_factor: 1
  separable: false
  spectral_layers: 3
  activation_function: gelu
  pos_embed: true
  dropout_filter: 0
  dropout_mlp: 0.1
  pos_emb_dropout: 0.0
  drop_path_rate: 0.1
  num_blocks: 8
  sparsity_threshold: 0.0
  normalization_layer: instance_norm
  hard_thresholding_fraction: 1.0
  checkpointing: 0
  time_scale_shift_before_filter: true
  loss_function:
    _target_: src.losses.losses.LpLoss
    p: 2
    relative: true
  name: SFNO
  verbose: true
  num_conditional_channels: 0
  with_time_emb: true
datamodule:
  data_dir: /data
  batch_size: 6
  eval_batch_size: 2
  num_workers: 8
  pin_memory: false
  persistent_workers: true
  drop_last: false
  shuffle_train_data: true
  verbose: true
  window: 1
  horizon: 6
  max_train_samples: null
  max_val_samples: 80
  forcing_names:
  - DSWRFtoa
  in_names:
  - HGTsfc
  - PRESsfc
  - surface_temperature
  - air_temperature_0
  - air_temperature_1
  - air_temperature_2
  - air_temperature_3
  - air_temperature_4
  - air_temperature_5
  - air_temperature_6
  - air_temperature_7
  - specific_total_water_0
  - specific_total_water_1
  - specific_total_water_2
  - specific_total_water_3
  - specific_total_water_4
  - specific_total_water_5
  - specific_total_water_6
  - specific_total_water_7
  - eastward_wind_0
  - eastward_wind_1
  - eastward_wind_2
  - eastward_wind_3
  - eastward_wind_4
  - eastward_wind_5
  - eastward_wind_6
  - eastward_wind_7
  - northward_wind_0
  - northward_wind_1
  - northward_wind_2
  - northward_wind_3
  - northward_wind_4
  - northward_wind_5
  - northward_wind_6
  - northward_wind_7
  out_names:
  - PRESsfc
  - surface_temperature
  - air_temperature_0
  - air_temperature_1
  - air_temperature_2
  - air_temperature_3
  - air_temperature_4
  - air_temperature_5
  - air_temperature_6
  - air_temperature_7
  - specific_total_water_0
  - specific_total_water_1
  - specific_total_water_2
  - specific_total_water_3
  - specific_total_water_4
  - specific_total_water_5
  - specific_total_water_6
  - specific_total_water_7
  - eastward_wind_0
  - eastward_wind_1
  - eastward_wind_2
  - eastward_wind_3
  - eastward_wind_4
  - eastward_wind_5
  - eastward_wind_6
  - eastward_wind_7
  - northward_wind_0
  - northward_wind_1
  - northward_wind_2
  - northward_wind_3
  - northward_wind_4
  - northward_wind_5
  - northward_wind_6
  - northward_wind_7
  auxiliary_names: null
  prediction_horizon: null
  _target_: src.datamodules.fv3gfs_ensemble.FV3GFSEnsembleDataModule
  data_dir_stats: /data/stats
  training_sub_paths: null
  prescriber:
    _target_: fme.core.prescriber.Prescriber
    prescribed_name: surface_temperature
    mask_name: ocean_fraction
    mask_value: 1
module:
  optimizer:
    name: adamw
    lr: 0.0004
    weight_decay: 1.0e-05
    eps: 1.0e-08
    betas:
    - 0.9
    - 0.99
  scheduler:
    _target_: torch.optim.lr_scheduler.CosineAnnealingLR
    T_max: 60
  monitor: val/avg/crps
  mode: min
  name: ''
  use_ema: true
  ema_decay: 0.9999
  enable_inference_dropout: true
  num_predictions: 16
  prediction_inputs_noise: 0.0
  logging_infix: ''
  log_every_step_up_to: 1000
  verbose: true
  seed: 11
  work_dir: ./results
  _target_: src.experiment_types.interpolation.InterpolationExperiment
  stack_window_to_channel_dim: true
callbacks:
  model_checkpoint:
    _target_: pytorch_lightning.callbacks.ModelCheckpoint
    monitor: val/avg/crps
    mode: min
    save_top_k: 1
    save_last: true
    verbose: true
    dirpath: ./results/checkpoints/01HER6GMAMHXPA62240XSV4TD4
    filename: FV3GFS-Ipol6h_None_epoch{epoch:03d}_seed11
    auto_insert_metric_name: false
  watch_model:
    _target_: src.utilities.wandb_callbacks.WatchModel
    log: all
    log_freq: 500
  summarize_best_val_metric:
    _target_: src.utilities.wandb_callbacks.SummarizeBestValMetric
  learning_rate_logging:
    _target_: pytorch_lightning.callbacks.LearningRateMonitor
logger:
  wandb:
    _target_: src.utilities.wandb_callbacks.MyWandbLogger
    entity: ai2cm
    name: FV3GFS-Ipol6h_SFNO_EMA_256x8h_L2R_44lr_10mlpDr_10dpr_15wd_cos_11seed
    tags:
    - fv3gfs
    - interpolation
    notes: '...'
    project: FME-v2-salva
    group: FV3GFS-Ipol6h_SFNO_EMA_256x8h_L2R_44lr_10mlpDr_10dpr_15wd_cos
    resume: allow
    reinit: true
    mode: online
    save_dir: ./results/
    offline: false
    id: 01HER6GMAMHXPA62240XSV4TD4
    log_model: false
    prefix: ''
n_gpus: 4
effective_batch_size: 72