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