From 0ca64a6687e73926bcafb514ae16fb0d4053dbc4 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Fri, 21 Nov 2025 08:33:44 +0100 Subject: [PATCH] Streamline model for better training --- Code/python/src/aiRNN/models.py | 25 ++++++++++--------------- 1 file changed, 10 insertions(+), 15 deletions(-) diff --git a/Code/python/src/aiRNN/models.py b/Code/python/src/aiRNN/models.py index 39ab19d..891ffbb 100644 --- a/Code/python/src/aiRNN/models.py +++ b/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