Streamline model for better training

Code/python/src/aiRNN/models.py

@@ -15,24 +15,24 @@ class BaseRNN(nn.Module):
 
         if rnn_type == "RNN":
             self.rnn = nn.RNN(
-                input_size=hidden_size,
+                input_size=hidden_size*3,
                 hidden_size=rnn_size,
+                num_layers=2,
                 batch_first=True,
                 device=device
             )
         elif rnn_type == "LSTM":
             self.rnn = nn.LSTM(
-                input_size=hidden_size,
+                input_size=hidden_size*3,
                 hidden_size=rnn_size,
+                num_layers=2,
                 batch_first=True,
                 device=device
             )
         else:
             raise ValueError("Unsupported rnn_type")
 
-        readout_in = rnn_size + time_in
-        if context_in is not None:
-            readout_in += context_in
+        readout_in = rnn_size
         self.readout = nn.Linear(readout_in, out_size, device=device)
 
         self.context_in = context_in
@@ -54,13 +54,13 @@ class BaseRNN(nn.Module):
             ctx = torch.zeros(b, total_len, time_proj.size(-1),
                                device=time_seq.device)
 
-        # Build the RNN input sequence
-        rnn_input = time_proj + ctx
-
+        feat_proj_or_zeros = torch.zeros(b, total_len, time_proj.size(-1), device=time_seq.device)
         if feat_init is not None:
             feat_proj = self.feat_proj(feat_init)               # (b, init_steps, hidden)
             # Insert feat input *only* in first init_steps
-            rnn_input[:, :init_steps, :] += feat_proj
+            feat_proj_or_zeros[:, :init_steps, :] += feat_proj
+        
+        rnn_input = torch.cat([time_proj, ctx, feat_proj_or_zeros], dim=-1)
 
         # Run the RNN over full sequence
         outputs, hidden_next = self.rnn(rnn_input, hidden)
@@ -71,12 +71,7 @@ class BaseRNN(nn.Module):
         # Build readout input
         time_raw = time_seq[:, init_steps:init_steps + pred_steps]
 
-        read_list = [pred_h, time_raw]
-        if context is not None:
-            read_list.append(context.unsqueeze(1).expand(-1, pred_steps, -1))
-        read_in = torch.cat(read_list, dim=-1)
-
-        preds = self.readout(read_in)
+        preds = self.readout(pred_h)
         return preds, hidden_next