Repository: AIStream-Peelout/flow-forecast

Issue:
Add vanilla GRU
We want to a simple vanilla GRU for time series forecasting maybe with the probabilistic component as well.


Discussion:

Resolving patch:
diff --git a/flood_forecast/basic/gru_vanilla.py b/flood_forecast/basic/gru_vanilla.py
new file mode 100755
--- /dev/null
+++ b/flood_forecast/basic/gru_vanilla.py
@@ -0,0 +1,55 @@
+import torch
+
+
+class VanillaGRU(torch.nn.Module):
+    def __init__(self, n_time_series: int, hidden_dim: int, num_layers: int, n_target: int, dropout: float,
+                 forecast_length=1, use_hidden=False):
+        """
+        Simple GRU to preform deep time series forecasting.
+
+        :param n_time_series:
+        :type n_time_series:
+        :param hidden_dim:
+        :type hidden_dim:
+        """
+        super(VanillaGRU, self).__init__()
+
+        # Defining the number of layers and the nodes in each layer
+        self.layer_dim = num_layers
+        self.hidden_dim = hidden_dim
+        self.hidden = None
+        self.use_hidden = use_hidden
+        self.forecast_length = forecast_length
+
+        # GRU layers
+        self.gru = torch.nn.GRU(
+            n_time_series, hidden_dim, num_layers, batch_first=True, dropout=dropout
+        )
+
+        # Fully connected layer
+        self.fc = torch.nn.Linear(hidden_dim, n_target)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward function for GRU
+
+        :param x: torch of shape
+        :type model: torch.Tensor
+        :return: Returns a tensor of shape (batch_size, forecast_length, n_target) or (batch_size, n_target)
+        :rtype: torch.Tensor
+        """
+        # Initializing hidden state for first input with zeros
+        if self.hidden is not None and self.use_hidden:
+            h0 = self.hidden
+        else:
+            h0 = torch.zeros(self.layer_dim, x.size(0), self.hidden_dim).requires_grad_()
+
+        # Forward propagation by passing in the input and hidden state into the model
+        out, self.hidden = self.gru(x, h0.detach())
+
+        # Reshaping the outputs in the shape of (batch_size, seq_length, hidden_size)
+        # so that it can fit into the fully connected layer
+        out = out[:, -self.forecast_length, :]
+
+        # Convert the final state to our desired output shape (batch_size, output_dim)
+        out = self.fc(out)
+        return out
diff --git a/flood_forecast/model_dict_function.py b/flood_forecast/model_dict_function.py
--- a/flood_forecast/model_dict_function.py
+++ b/flood_forecast/model_dict_function.py
@@ -18,6 +18,7 @@
 from flood_forecast.custom.dilate_loss import DilateLoss
 from flood_forecast.meta_models.basic_ae import AE
 from flood_forecast.transformer_xl.dsanet import DSANet
+from flood_forecast.basic.gru_vanilla import VanillaGRU
 
 """
 Utility dictionaries to map a string to a class
@@ -34,7 +35,8 @@
     "DecoderTransformer": DecoderTransformer,
     "BasicAE": AE,
     "Informer": Informer,
-    "DSANet": DSANet
+    "DSANet": DSANet,
+    "VanillaGRU": VanillaGRU
 }
 
 pytorch_criterion_dict = {


Test patch:
diff --git a/tests/gru_vanilla.json b/tests/gru_vanilla.json
new file mode 100644
--- /dev/null
+++ b/tests/gru_vanilla.json
@@ -0,0 +1,75 @@
+{                  
+    "model_name": "VanillaGRU",
+    "use_decoder": true,
+    "model_type": "PyTorch",
+    "model_params": {
+      "n_time_series":3,
+      "hidden_dim": 20,
+      "num_layers": 1,
+      "n_target": 1,
+      "dropout": 0.11
+   },
+   "dataset_params":
+   {  "class": "default",
+      "training_path": "tests/test_data/keag_small.csv",
+      "validation_path": "tests/test_data/keag_small.csv",
+      "test_path": "tests/test_data/keag_small.csv",
+      "batch_size":10,
+      "forecast_history":10,
+      "forecast_length":1,
+      "train_start": 1,
+      "train_end": 300,
+      "valid_start":302,
+      "valid_end": 401,
+      "test_start":50,
+      "test_end": 450,
+      "target_col": ["cfs"],
+      "relevant_cols": ["cfs", "precip", "temp"],
+      "scaler": "StandardScaler", 
+      "interpolate": false
+   },
+   "training_params":
+   {
+      "criterion":"MSE",
+      "optimizer": "Adam",
+      "optim_params":
+      {
+
+      },
+      "lr": 0.03,
+      "epochs": 1,
+      "batch_size":4
+   
+   },
+   "GCS": false,
+   
+   "wandb": {
+      "name": "flood_forecast_circleci",
+      "tags": ["dummy_run", "circleci"],
+      "project": "repo-flood_forecast"
+   },
+   "forward_params":{},
+  "metrics":["MSE"],
+  "inference_params":
+  {     
+        "datetime_start":"2016-05-31",
+         "hours_to_forecast":334, 
+         "test_csv_path":"tests/test_data/keag_small.csv",
+         "decoder_params":{
+           "decoder_function": "simple_decode", 
+           "unsqueeze_dim": 1}
+        ,
+         "dataset_params":{
+            "file_path": "tests/test_data/keag_small.csv",
+            "forecast_history":10,
+            "forecast_length":1,
+            "relevant_cols": ["cfs", "precip", "temp"],
+            "target_col": ["cfs"],
+            "scaling": "StandardScaler",
+            "interpolate_param": false
+         }
+  }
+}
+      
+  
+   
\ No newline at end of file