Revert "CorrDiff: CWB Dataloader cleanups (#593)" (#626)

This reverts commit 4525f08.
NVIDIA · Jul 29, 2024 · 9994b14 · 9994b14
1 parent edccbb3
commit 9994b14
Show file tree

Hide file tree

Showing 8 changed files with 343 additions and 95 deletions.
diff --git a/examples/generative/corrdiff/conf/config_generate.yaml b/examples/generative/corrdiff/conf/config_generate.yaml
@@ -82,16 +82,19 @@ dataset:
   train: False
   all_times: True
   data_path: /data/data2/2023-01-24-cwb-4years.zarr
+  n_history: 0
   in_channels: [0, 1, 2, 3, 4, 9, 10, 11, 12, 17, 18, 19]
   out_channels: [0, 17, 18, 19]
   img_shape_x: 448
   img_shape_y: 448
+  roll: false
   add_grid: true
   ds_factor: 4
   min_path: null
   max_path: null
   global_means_path: null
   global_stds_path: null
+  normalization: "v1"
 
 ## Deterministic sampler options
 sigma_min: null

diff --git a/examples/generative/corrdiff/conf/config_generate_patch.yaml b/examples/generative/corrdiff/conf/config_generate_patch.yaml
@@ -82,16 +82,19 @@ dataset:
   train: False
   all_times: True
   data_path: /data/data2/2023-01-24-cwb-4years.zarr
+  n_history: 0
   in_channels: [0, 1, 2, 3, 4, 9, 10, 11, 12, 17, 18, 19]
   out_channels: [0, 17, 18, 19]
   img_shape_x: 448
   img_shape_y: 448
+  roll: false
   add_grid: true
   ds_factor: 4
   min_path: null
   max_path: null
   global_means_path: null
   global_stds_path: null
+  normalization: "v1"
 
 
 ## Deterministic sampler options

diff --git a/examples/generative/corrdiff/conf/config_train_base.yaml b/examples/generative/corrdiff/conf/config_train_base.yaml
@@ -87,16 +87,19 @@ dataset:
   type: cwb
   train_test_split: false
   data_path: /data/data2/2023-01-24-cwb-4years.zarr
+  n_history: 0
   in_channels: [0, 1, 2, 3, 4, 9, 10, 11, 12, 17, 18, 19]
   out_channels: [0, 17, 18, 19]
   img_shape_x: 448
   img_shape_y: 448
+  roll: false
   add_grid: true
   ds_factor: 4
   min_path: null
   max_path: null
   global_means_path: null
   global_stds_path: null
+  normalization: "v1"
 
 # Validation dataset options
 # (need to set dataset.train_test_split == true to have an effect)

diff --git a/examples/generative/corrdiff/conf/references/config_data_ref.yaml b/examples/generative/corrdiff/conf/references/config_data_ref.yaml
@@ -18,8 +18,10 @@ full_field: &FULL_FIELD
 
   batch_size: 1
   dt: 1
+  n_history: 0
   img_shape_x: 448
   img_shape_y: 448
+  normalization: v1     #minmax
   # era5_data_dir: '/lustre/fsw/sw_climate_fno/cwb-align'  #old
   # cwb_data_dir: '/lustre/fsw/sw_climate_fno/cwb-rwrf-pad-2212/all_ranges'  #old
   # relative to CWB_ROOT environment variable
@@ -38,6 +40,7 @@ full_field_train_crop448_grid_12inchans_fcn_4outchans_4x_6layer: #config for sin
   <<: *FULL_FIELD
   in_channels: [0,1,2,3,4,9,10,11,12,17,18,19] 
   out_channels: [0,17,18,19]   #[17,18,19]
+  roll: !!bool False
   patch_size: 448
   crop_size_x: 448
   crop_size_y: 448
@@ -50,6 +53,7 @@ full_field_val_crop448_grid_12inchans_fcn_4outchans_4x_6layer: #config for singl
   <<: *FULL_FIELD
   in_channels: [0,1,2,3,4,9,10,11,12,17,18,19]
   out_channels: [0,17,18,19]  #[0,9,10,11] 
+  roll: !!bool False
   patch_size: 448
   crop_size_x: 448
   crop_size_y: 448
@@ -61,6 +65,7 @@ out_of_sample:
   <<: *FULL_FIELD
   in_channels: [0,1,2,3,4,9,10,11,12,17,18,19]
   out_channels: [0,17,18,19]  #[0,9,10,11] 
+  roll: !!bool False
   patch_size: 448
   crop_size_x: 448
   crop_size_y: 448
@@ -76,12 +81,14 @@ full_field_train_crop448_grid_12inchans_fcn_4outchans_4x_normv2: #config for sin
   <<: *FULL_FIELD
   in_channels: [0,1,2,3,4,9,10,11,12,17,18,19] 
   out_channels: [0,17,18,19]   #[17,18,19]
+  roll: !!bool False
   patch_size: 448
   crop_size_x: 448
   crop_size_y: 448
   N_grid_channels: 4
   add_grid: True
   ds_factor: 4
+  normalization: v2
 
 
 
@@ -90,6 +97,7 @@ full_field_train_crop448_grid_12inchans_fcn_4outchans_4x: #config for single gpu
   <<: *FULL_FIELD
   in_channels: [0,1,2,3,4,9,10,11,12,17,18,19] 
   out_channels: [0,17,18,19]   #[17,18,19]
+  roll: !!bool False
   patch_size: 448
   crop_size_x: 448
   crop_size_y: 448
@@ -102,6 +110,7 @@ full_field_val_crop448_grid_12inchans_fcn_4outchans_4x:  &validation #config for
   <<: *FULL_FIELD
   in_channels: [0,1,2,3,4,9,10,11,12,17,18,19]
   out_channels: [0,17,18,19]  #[0,9,10,11] 
+  roll: !!bool False
   patch_size: 448
   crop_size_x: 448
   crop_size_y: 448
@@ -115,6 +124,7 @@ full_field_train_crop448_grid_20inchans_4outchans_4x: #config for single gpu tra
   <<: *FULL_FIELD
   in_channels: [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
   out_channels: [0,17,18,19]   #[17,18,19]
+  roll: !!bool False
   patch_size: 448
   crop_size_x: 448
   crop_size_y: 448
@@ -127,6 +137,7 @@ full_field_val_crop448_grid_20inchans_4outchans_4x: #config for single gpu train
   <<: *FULL_FIELD
   in_channels: [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
   out_channels: [0,17,18,19]   #[17, 18, 19]
+  roll: !!bool False
   patch_size: 448
   crop_size_x: 448
   crop_size_y: 448
@@ -140,6 +151,7 @@ full_field_train_crop112_grid_20inchans_4outchans_4x: #config for single gpu tra
   <<: *FULL_FIELD
   in_channels: [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
   out_channels: [0,17,18,19]   #[17,18,19]
+  roll: !!bool False
   patch_size: 112
   crop_size_x: 112
   crop_size_y: 112
@@ -152,6 +164,7 @@ full_field_val_crop112_grid_20inchans_4outchans_4x: #config for single gpu train
   <<: *FULL_FIELD
   in_channels: [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
   out_channels: [0,17,18,19]   #[17, 18, 19]
+  roll: !!bool False
   patch_size: 112
   crop_size_x: 448
   crop_size_y: 448
@@ -165,6 +178,7 @@ full_field_train_crop112_grid_20inchans_19outchans_4x: #config for single gpu tr
   <<: *FULL_FIELD
   in_channels: [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
   out_channels: [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]   #[17,18,19]
+  roll: !!bool False
   patch_size: 112
   crop_size_x: 112
   crop_size_y: 112
@@ -177,6 +191,7 @@ full_field_val_crop112_grid_20inchans_19outchans_4x: #config for single gpu trai
   <<: *FULL_FIELD
   in_channels: [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
   out_channels: [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]   #[17, 18, 19]
+  roll: !!bool False
   patch_size: 112
   crop_size_x: 448
   crop_size_y: 448
@@ -191,6 +206,7 @@ full_field_train_crop112_grid_20inchans_4x: #config for single gpu training to c
   <<: *FULL_FIELD
   in_channels: [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
   out_channels: [17,18,19]  #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
+  roll: !!bool False
   patch_size: 112
   crop_size_x: 112
   crop_size_y: 112
@@ -203,6 +219,7 @@ full_field_val_crop112_grid_20inchans_4x: #config for single gpu training to cat
   <<: *FULL_FIELD
   in_channels: [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
   out_channels: [17, 18, 19]     #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
+  roll: !!bool False
   patch_size: 112
   crop_size_x: 448
   crop_size_y: 448
@@ -216,6 +233,7 @@ full_field_train_crop64_grid_20inchans_4x: #config for single gpu training to ca
   <<: *FULL_FIELD
   in_channels: [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
   out_channels: [17, 18, 19]     #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
+  roll: !!bool False
   patch_size: 64
   crop_size_x: 64
   crop_size_y: 64
@@ -228,6 +246,7 @@ full_field_val_crop64_grid_20inchans_4x: #config for single gpu training to catc
   <<: *FULL_FIELD
   in_channels: [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
   out_channels: [17, 18, 19]     #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
+  roll: !!bool False
   patch_size: 64
   crop_size_x: 448
   crop_size_y: 448
@@ -240,6 +259,7 @@ full_field_train_crop64_grid_20inchans: #config for single gpu training to catch
   <<: *FULL_FIELD
   in_channels: [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
   out_channels: [17, 18, 19]     #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
+  roll: !!bool False
   patch_size: 64
   crop_size_x: 64
   crop_size_y: 64
@@ -251,6 +271,7 @@ full_field_val_crop64_grid_20inchans: #config for single gpu training to catch b
   <<: *FULL_FIELD
   in_channels: [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
   out_channels: [17, 18, 19]     #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
+  roll: !!bool False
   patch_size: 64
   crop_size_x: 448
   crop_size_y: 448
@@ -263,6 +284,7 @@ full_field_train_crop64_grid: #config for single gpu training to catch bugs and
   <<: *FULL_FIELD
   in_channels: [17, 18, 19]  #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
   out_channels: [17, 18, 19]     #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
+  roll: !!bool False
   patch_size: 64
   crop_size_x: 64
   crop_size_y: 64
@@ -274,6 +296,7 @@ full_field_val_crop64_grid: #config for single gpu training to catch bugs and ov
   <<: *FULL_FIELD
   in_channels: [17, 18, 19]      #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
   out_channels: [17, 18, 19]     #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
+  roll: !!bool False
   patch_size: 64
   crop_size_x: 448
   crop_size_y: 448
@@ -287,6 +310,7 @@ full_field_train_crop64: #config for single gpu training to catch bugs and overf
   <<: *FULL_FIELD
   in_channels: [17.18,19]        #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
   out_channels: [17, 18, 19]     #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
+  roll: !!bool False
   patch_size: 64
   crop_size_x: 64
   crop_size_y: 64
@@ -296,6 +320,7 @@ full_field_val_crop64: #config for single gpu training to catch bugs and overfit
   <<: *FULL_FIELD
   in_channels: [17, 18, 19]      #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
   out_channels: [17, 18, 19]     #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
+  roll: !!bool False
   patch_size: 64
   crop_size_x: 448
   crop_size_y: 448
@@ -306,6 +331,7 @@ full_field_train_crop112: #config for single gpu training to catch bugs and over
   <<: *FULL_FIELD
   in_channels: [17.18,19]        #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
   out_channels: [17, 18, 19]     #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
+  roll: !!bool False
   patch_size: 112
   crop_size_x: 112
   crop_size_y: 112
@@ -315,6 +341,7 @@ full_field_val_crop112: #config for single gpu training to catch bugs and overfi
   <<: *FULL_FIELD
   in_channels: [17, 18, 19]      #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
   out_channels: [17, 18, 19]     #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
+  roll: !!bool False
   patch_size: 112
   crop_size_x: 112
   crop_size_y: 112
@@ -325,6 +352,7 @@ full_field_train_crop224: #config for single gpu training to catch bugs and over
   <<: *FULL_FIELD
   in_channels: [17.18,19]        #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
   out_channels: [17, 18, 19]     #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
+  roll: !!bool False
   patch_size: 224
   crop_size_x: 224
   crop_size_y: 224
@@ -334,6 +362,7 @@ full_field_val_crop224: #config for single gpu training to catch bugs and overfi
   <<: *FULL_FIELD
   in_channels: [17, 18, 19]      #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
   out_channels: [17, 18, 19]     #[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19]
+  roll: !!bool False
   patch_size: 224
   crop_size_x: 224
   crop_size_y: 224