[Diff] sbAsma/issue1279 noise conditioning

config/default_config.yml

@@ -42,16 +42,15 @@ pred_mlp_adaln: True
 
 # number of steps offset applied to first target window; if set to zero and forecast_steps=0 then
 # one is training an auto-encoder
-forecast_offset : 1
+forecast_offset : 0
 forecast_delta_hrs: 0
-forecast_steps: 1
-forecast_policy: "fixed"
+forecast_steps: 0
+forecast_policy: null
 forecast_att_dense_rate: 1.0
-fe_num_blocks: 8
+fe_num_blocks: 0
 fe_num_heads: 16
 fe_dropout_rate: 0.1
 fe_with_qk_lnorm: True
-fe_diffusion_model: True
 impute_latent_noise_std: 0.0  # 1e-4
 chkpt_encoder_weights: "./models/whkujigw/whkujigw_epoch00063.chkpt"
 
@@ -85,28 +84,14 @@ validate_with_ema: True
 ema_ramp_up_ratio: 0.09
 ema_halflife_in_thousands: 1e-3
 
-# training mode: "forecast" or "masking" (masked token modeling) or "student-teacher"
+# training mode: "forecast" or "masking" (masked token modeling)
 # for "masking" to train with auto-encoder mode, forecast_offset should be 0
-training_mode: "forecast"
-"target_and_aux_calc": "DiffusionLatentTargetEncoder"
-training_mode_config: {
-  "losses" : { 
-    LossLatentDiffusion: {
-        weight: 1.0, loss_fcts: [['mse', 1.0]]
-    }
-  },
-  "shared_heads": False,
-  "teacher_model": {}
-}
-validation_mode_config: {
-  "losses" : { 
-    LossLatentDiffusion: {
-        weight: 1.0, loss_fcts: [['mse', 1.0]]
-    }
-  },
-  "shared_heads": False,
-  "teacher_model": {}
-}
+training_mode: "masking"
+training_mode_config: {"losses": {LossPhysical: {weight: 1.0, loss_fcts: [['mse', 1.0]]},
+                                  }
+                      }
+validation_mode_config: {"losses": {LossPhysical: {weight: 1.0, loss_fcts: [['mse', 1.0]]},}
+                        }
 # masking rate when training mode is "masking"; ignored in foreacast mode
 masking_rate: 0.6
 # sample the masking rate (with normal distribution centered at masking_rate)

config/diffusion_config.yml

@@ -0,0 +1,52 @@
+# Changes from original to new config
+forecast_offset: 1  # changed from 0
+forecast_steps: 1   # changed from 0
+forecast_policy: "fixed"  # changed from null
+fe_num_blocks: 8    # changed from 0
+fe_diffusion_model: True  # new parameter added
+
+training_mode: "forecast"
+"target_and_aux_calc": "DiffusionLatentTargetEncoder"
+training_mode_config: {
+  "losses" : { 
+    LossLatentDiffusion: {
+        weight: 1.0, loss_fcts: [['mse', 1.0]]
+    }
+  },
+  "shared_heads": False,
+  "teacher_model": {}
+}
+validation_mode_config: {
+  "losses" : { 
+    LossLatentDiffusion: {
+        weight: 1.0, loss_fcts: [['mse', 1.0]]
+    }
+  },
+  "shared_heads": False,
+  "teacher_model": {}
+}
+
+# training_mode: "forecast"  # changed from "masking"
+# training_mode_config:
+#   losses:
+#     LossLatentDiffusion:  # changed from LossPhysical
+#       weight: 1.0
+#       loss_fcts: [['mse', 1.0]]
+#   shared_heads: False  # new parameter
+#   target_and_aux_calc: "DiffusionLatentTargetEncoder"  # new parameter
+#   teacher_model: {}  # new parameter
+
+# validation_mode_config:
+#   losses:
+#     LossLatentDiffusion:  # changed from LossPhysical
+#       weight: 1.0
+#       loss_fcts: [['mse', 1.0]]
+#   shared_heads: False  # new parameter
+#   target_and_aux_calc: "DiffusionLatentTargetEncoder"  # new parameter
+#   teacher_model: {}  # new parameter
+
+samples_per_mini_epoch: 32  # changed from 4096
+samples_per_validation: 8   # changed from 512
+
+# Updated comment for training_mode
+# training mode: "forecast" or "masking" (masked token modeling) or "student-teacher"

src/weathergen/model/model.py

@@ -332,9 +332,15 @@ def create(self) -> "Model":
                 "Empty forecast engine (fe_num_blocks = 0), but forecast_steps[i] > 0 for some i"
             )
 
-        self.forecast_engine = ForecastingEngine(cf, self.num_healpix_cells)
-        if cf.fe_diffusion_model:
-            self.forecast_engine = DiffusionForecastEngine(forecast_engine=self.forecast_engine)
+        # check if diffusion mode is enabled
+        fe_diffusion_model = getattr(cf, "fe_diffusion_model", False)
+        if fe_diffusion_model:
+            self.forecast_engine = DiffusionForecastEngine(
+                    forecast_engine=ForecastingEngine(cf, self.num_healpix_cells)
+                )
+        else:
+            self.forecast_engine = ForecastingEngine(cf, self.num_healpix_cells)
+
 
         ###############
         # embed coordinates yielding one query token for each target token