Add small UNET model

ECCC/configs/config_UNET_small.yaml

@@ -0,0 +1,74 @@
+data:
+  type: eccc
+  data_training_index: ./granite-geospatial-wxc-downscaling/ECCC/indices/index_example.json
+  data_val_index: ./granite-geospatial-wxc-downscaling/ECCC/indices/index_example.json
+  data_test_index: ./granite-geospatial-wxc-downscaling/ECCC/indices/index_example.json
+  static_data_index: ./granite-geospatial-wxc-downscaling/ECCC/indices/staticDataIndex.json
+  input_surface_vars: [TT, UUWE, VVSN]
+  input_static_surface_vars: [ME, MG, Z0]
+  vertical_pres_vars: [GZ_pressure_levels, HU_pressure_levels, TT_pressure_levels,
+    UUWE_pressure_levels, VVSN_pressure_levels, WW_pressure_levels]
+  vertical_level1_vars: [GZ_model_levels, HU_model_levels, TT_model_levels, WW_model_levels]
+  vertical_level2_vars: [UUWE_model_levels, VVSN_model_levels]
+  other: [NT, P0, PC, PR, TM, H]
+  input_level_pres: [500.0, 700.0, 850.0, 925.0, 1000.0]
+  input_level1: [0.989269, 0.99501, 0.99875]
+  input_level2: [0.986026, 0.992524, 0.997502]
+  output_vars: [UUWE, VVSN]
+  n_input_timestamps: 1
+  input_size_lat: 32
+  input_size_lon: 32
+  apply_smoothen: false
+  downsample_factor: 8
+  n_random_windows: 16
+  crop_factor: 256
+model:
+  input_mu: ./granite-geospatial-wxc-downscaling/ECCC/scalers/input_mu.pt
+  input_sigma: ./granite-geospatial-wxc-downscaling/ECCC/scalers/input_sigma.pt
+  input_static_mu: ./granite-geospatial-wxc-downscaling/ECCC/scalers/input_static_mu.pt
+  input_static_sigma: ./granite-geospatial-wxc-downscaling/ECCC/scalers/input_static_sigma.pt
+  target_mu: ./granite-geospatial-wxc-downscaling/ECCC/scalers/target_mu.pt
+  target_sigma: ./granite-geospatial-wxc-downscaling/ECCC/scalers/target_sigma.pt
+  target_static_mu: ./granite-geospatial-wxc-downscaling/ECCC/scalers/target_static_mu.pt
+  target_static_sigma: ./granite-geospatial-wxc-downscaling/ECCC/scalers/target_static_sigma.pt
+  embed_dim: 1024
+  n_blocks_encoder: 8
+  mlp_multiplier: 4
+  n_heads: 16
+  dropout_rate: 0.0
+  drop_path: 0.0
+  residual: none
+  num_static_channels: 7
+  downscaling_patch_size: [3, 3]
+  downscaling_embed_dim: 512
+  encoder_decoder_type: conv
+  encoder_decoder_upsampling_mode: pixel_shuffle
+  encoder_decoder_kernel_size_per_stage: [[]]
+  encoder_decoder_scale_per_stage: [[]]
+  encoder_decoder_conv_channels: 128
+  residual_connection: true
+  encoder_shift: false
+  unet: true
+job_id: example
+type_training: one_per_epoch
+batch_size: 1
+num_epochs: 10
+dl_num_workers: 1
+dl_prefetch_size: 1
+learning_rate: 1.0e-05
+limit_steps_train: 832
+limit_steps_valid: 20
+min_lr: 1.0e-06
+max_lr: 5.0e-05
+warm_up_steps: 0
+mask_unit_size: [2, 2]
+mask_ratio_inputs: 0.0
+mask_ratio_targets: 0.0
+max_batch_size: 16
+path_experiment: ./granite-geospatial-wxc-downscaling/ECCC
+path_model_weights: ./granite-geospatial-wxc-downscaling/ECCC/weights/best_rmse_UNET_small.pt
+backbone_use: true
+backbone_freeze: false
+backbone_prefix: encoder.
+finetune_w_static: true
+strict_matching: false

ECCC/weights/best_rmse_UNET_small.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:16fe7b3fa6b3724c0601d051b712be0345de5e53b4a4af6e08f7c5969f899fac
+size 3056019514