File size: 5,030 Bytes
407ccb1 | 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 | {
"experiment": {
"name": "pixeldit_sem1024d24l_pix128d4l_bilineardwgeluproject_unfreeze_3src_const1em5_t128_4stepft",
"rundir": "",
"model_type": "pixeldit",
"resume_from": "",
"reset_best_valid_loss": true,
"reset_scheduler_state": true,
"skip_optimizer_reloading": true
},
"dit": {
"patch_size": 1,
"patch_size_vert": 1,
"patch_size_horiz": 4,
"embed_dim": 1024,
"n_layers": 24,
"num_heads": 16,
"qk_norm": true,
"qk_norm_elementwise_affine": false,
"attn_kernel": 9,
"do_interleaved_dilation": false,
"na_dilations": 3,
"na3d_backend": "",
"gated_attention": true,
"do_alt_depthwise_attn": true,
"use_conv_instead_of_na3d": false,
"conv_kernel_size": 3,
"do_rope_2d": false,
"do_rope_2d_stereographic": true,
"do_rotate_wind": true,
"do_absolute_pos_embedding": false,
"index_is_latlon": true,
"do_arcsinh_input": false,
"do_te_mlp_last_layer": false,
"use_cross_layer_concat_skip_connection": false,
"do_activation_checkpointing": 1.0
},
"pixel_dit": {
"embed_dim": 128,
"n_layers": 4,
"num_heads": null,
"mlp_ratio": 4.0,
"attn_kernel": 9,
"gated_attention": false,
"qk_norm": true,
"qk_norm_elementwise_affine": false,
"freeze_semantic": false,
"freeze_pixel_blocks": false,
"use_bilinear_dw_gelu_project_adaln": true,
"do_activation_checkpointing": 0.0
},
"pipeline": {
"target_type": "mixed_state_diff",
"do_qv_softplus": true,
"do_precip_relu": true,
"variables_input_zeroed": []
},
"compute": {
"do_bf16_mixed": true,
"do_torch_compile": true
},
"training": {
"lr": 1e-05,
"loss_lr": 0.1,
"epochs": 50,
"batch_size": 1,
"num_workers": 2,
"num_workers_inference": 2,
"optimizer": "adamw",
"beta1": 0.9,
"beta2": 0.999,
"weight_decay": 1e-05,
"seed": 42,
"scheduler_type": "constant",
"warmup_steps": 2000,
"cosine_t_max": 100,
"cosine_n_cycles": 5,
"cosine_decay": 0.3162,
"num_steps": 4,
"multistep_training_mode": "final_only",
"do_tf32": true,
"grad_clip_max_norm": 0.5,
"backup_every_steps": 100,
"loss_type": "smooth_l1",
"resume_from": null,
"reset_best_valid_loss": false,
"reset_scheduler_state": false,
"skip_optimizer_reloading": false,
"fit_batches_for_norm": 1280,
"normalization_ref_path_input": "",
"normalization_ref_path_target": "",
"do_dual_normalization": false,
"pooled_loss": null,
"pooled_loss_channel_weights_csv": null,
"max_steps_per_epoch": null,
"validate_steps": 400,
"validate_min_samples": 1024,
"validate_total_images": null,
"validate_num_shards": null,
"exit_after_first_validation": false,
"max_validations": 5,
"no_save": false,
"save_ckpt_every_epoch": true,
"use_time_variation_seed": true,
"do_backtest_inference": false,
"do_track_gns": false,
"gns_decay": 0.99,
"gns_freq": 1
},
"data": {
"grid_type": "cubesphere",
"cubesphere_ne": 1024,
"cubesphere_npg": 2,
"tile_size": 128,
"nside": 2048,
"level_start": 8,
"level_end": 32,
"plevel": 1,
"cross_face_tiles": false,
"use_duo_padding": false,
"duo_padding_scrip_src_path": null,
"duo_padding_scrip_tgt_path": null,
"skip_corner_tiles": false,
"balance_cross_face": false,
"variables_prognostic": [
"PotentialTemperature",
"U",
"V",
"z_mid",
"omega",
"qv",
"T_2m"
],
"variables_prognostic_state": [
"omega"
],
"variables_forcing": [
"coszr",
"phis",
"sgh30",
"landfrac"
],
"variables_diagnostic": [
"precip_ice_surf_mass_flux",
"precip_liq_surf_mass_flux",
"ps"
],
"train_main_zarr_paths": [
"sdecadal.ne1024pg2_ne1024pg2.F20TR-SCREAMv1.c10-sep11.out10min.cubesphere.zarr",
"sdy2.ne1024pg2_ne1024pg2.F2010-SCREAMv1.c10-sep11-f602da2b98.out10min-4lev.cubesphere.zarr",
"sdy1.ne1024pg2_ne1024pg2.F2010-SCREAMv1.c10-sep11-f602da2b98.out10min-4lev.cubesphere.zarr"
],
"train_aux_zarr_paths": [
"sdecadal.ne1024pg2_ne1024pg2.F20TR-SCREAMv1.c10-sep11.out10min.cubesphere.zarr",
"sdy2.ne1024pg2_ne1024pg2.F2010-SCREAMv1.c10-sep11-f602da2b98.out10min-4lev.cubesphere.zarr",
"sdy1.ne1024pg2_ne1024pg2.F2010-SCREAMv1.c10-sep11-f602da2b98.out10min-4lev.cubesphere.zarr"
],
"train_zarr_weights": [
0.5,
1.0,
1.0
],
"train_start_indices": [
144,
144,
144
],
"train_end_indices": [
1435,
720,
720
],
"split_by_time": true,
"train_start_index": 144,
"train_end_index": 1435,
"test_start_index": 1872,
"test_end_index": 2015,
"test_stride": 36,
"inference_start_index": 1872,
"backtest_main_zarr_path": "",
"backtest_aux_zarr_path": "",
"backtest_start_index": 864,
"backtest_end_index": 1007,
"backtest_stride": 36,
"latlon_path": "latlon_ne1024pg2.nc",
"do_fixed_z": false,
"fixed_z_npy_path": "data/z_coordinate_24levels.npy",
"phis_npy_path": "data/phis.npy"
},
"debug": {
"print_steps": 5,
"mock": false,
"skip_india_ocean_eval": true,
"use_full_test_set": true,
"recreate_dataloader_each_validation": true,
"wandb_disabled": false
}
} |