Merge remote-tracking branch 'downstream/dev' into dev

Author: OuyangWenyu

torchhydro/models/model_dict_function.py

@@ -1,7 +1,7 @@
 """
 Author: Wenyu Ouyang
 Date: 2021-12-31 11:08:29
-LastEditTime: 2025-01-14 19:30:54
+LastEditTime: 2025-01-25 09:33:46
 LastEditors: Wenyu Ouyang
 Description: Dicts including models (which are seq-first), losses, and optims
 FilePath: /torchhydro/torchhydro/models/model_dict_function.py
@@ -19,6 +19,7 @@
 )
 
 from torchhydro.models.simple_lstm import (
+    LinearMultiLayerLSTMModel,
     LinearSimpleLSTMModel,
     MultiLayerLSTM,
     SimpleLSTM,
@@ -79,6 +80,7 @@
     "SimpleLSTM": SimpleLSTM,
     "LinearSimpleLSTMModel": LinearSimpleLSTMModel,
     "MultiLayerLSTM": MultiLayerLSTM,
+    "LinearMultiLayerLSTMModel": LinearMultiLayerLSTMModel,
 }
 
 pytorch_criterion_dict = {

torchhydro/models/simple_lstm.py

@@ -72,6 +72,25 @@ def forward(self, x):
         x0 = F.relu(self.former_linear(x))
         return super(LinearSimpleLSTMModel, self).forward(x0)
 
+class LinearMultiLayerLSTMModel(MultiLayerLSTM):
+    """
+    This model is nonlinear layer + MultiLayerLSTM.
+    """
+
+    def __init__(self, linear_size, **kwargs):
+        """
+
+        Parameters
+        ----------
+        linear_size
+            the number of input features for the first input linear layer
+        """
+        super(LinearMultiLayerLSTMModel, self).__init__(**kwargs)
+        self.former_linear = nn.Linear(linear_size, kwargs["input_size"])
+
+    def forward(self, x):
+        x0 = F.relu(self.former_linear(x))
+        return super(LinearMultiLayerLSTMModel, self).forward(x0)
 
 class SimpleLSTMForecast(SimpleLSTM):
     def __init__(self, input_size, output_size, hidden_size, forecast_length, dr=0.0):

torchhydro/trainers/deep_hydro.py

@@ -737,7 +737,12 @@ def load_model(self, mode="train"):
             raise NotImplementedError(
                 "For transfer learning, we need a pre-trained model"
             )
-        model = super().load_model(mode)
+        if mode == "train":
+            model = super().load_model(mode)
+        elif mode == "infer":
+            self.weight_path = self._get_trained_model()
+            model = self._load_model_from_pth()
+            model.to(self.device)
         if (
             "weight_path_add" in model_cfgs
             and "freeze_params" in model_cfgs["weight_path_add"]