Merge branch 'dev' of github.com:OuyangWenyu/torchhydro into dev

Author: OuyangWenyu

torchhydro/models/simple_lstm.py

@@ -1,7 +1,7 @@
 """
 Author: Wenyu Ouyang
 Date: 2023-09-19 09:36:25
-LastEditTime: 2025-01-25 09:19:43
+LastEditTime: 2025-01-25 09:32:12
 LastEditors: Wenyu Ouyang
 Description: Some self-made LSTMs
 FilePath: /torchhydro/torchhydro/models/simple_lstm.py
@@ -68,11 +68,9 @@ def __init__(self, linear_size, **kwargs):
         super(LinearSimpleLSTMModel, self).__init__(**kwargs)
         self.former_linear = nn.Linear(linear_size, kwargs["input_size"])
 
-    def forward(self, x, do_drop_mc=False, dropout_false=False):
+    def forward(self, x):
         x0 = F.relu(self.former_linear(x))
-        return super(LinearSimpleLSTMModel, self).forward(
-            x0, do_drop_mc=do_drop_mc, dropout_false=dropout_false
-        )
+        return super(LinearSimpleLSTMModel, self).forward(x0)
 
 
 class SimpleLSTMForecast(SimpleLSTM):

torchhydro/trainers/resulter.py

@@ -102,8 +102,10 @@ def save_result(self, pred, obs):
         save_dir = self.result_dir
         flow_pred_file = os.path.join(save_dir, self.pred_name)
         flow_obs_file = os.path.join(save_dir, self.obs_name)
-        pred.to_netcdf(flow_pred_file + ".nc")
-        obs.to_netcdf(flow_obs_file + ".nc")
+        max_len = max(len(basin) for basin in pred.basin.values)
+        encoding = {"basin": {"dtype": f"U{max_len}"}}
+        pred.to_netcdf(flow_pred_file + ".nc", encoding=encoding)
+        obs.to_netcdf(flow_obs_file + ".nc", encoding=encoding)
 
     def eval_result(self, preds_xr, obss_xr):
         # types of observations

torchhydro/trainers/trainer.py

@@ -1,7 +1,7 @@
 """
 Author: Wenyu Ouyang
 Date: 2021-12-05 11:21:58
-LastEditTime: 2024-05-23 15:21:17
+LastEditTime: 2025-01-14 19:29:34
 LastEditors: Wenyu Ouyang
 Description: Main function for training and testing
 FilePath: \torchhydro\torchhydro\trainers\trainer.py
@@ -61,13 +61,19 @@ def train_and_evaluate(cfgs: Dict):
     set_random_seed(random_seed)
     resulter = Resulter(cfgs)
     deephydro = _get_deep_hydro(cfgs)
-    if cfgs["training_cfgs"]["train_mode"] and (
+    # if train_mode is False, we only evaluate the model
+    train_mode = deephydro.cfgs["training_cfgs"]["train_mode"]
+    # but if train_mode is True, we still need some conditions to train the model
+    continue_train = deephydro.cfgs["model_cfgs"]["continue_train"]
+    is_transfer_learning = deephydro.cfgs["model_cfgs"]["model_type"] == "TransLearn"
+    is_train = train_mode and (
         (
             deephydro.weight_path is not None
-            and deephydro.cfgs["model_cfgs"]["continue_train"]
+            and (continue_train or is_transfer_learning)
         )
         or (deephydro.weight_path is None)
-    ):
+    )
+    if is_train:
         deephydro.model_train()
     preds, obss = deephydro.model_evaluate()
     resulter.save_cfg(deephydro.cfgs)