multi-layer lstm

Author: sfluegel05

chebai/models/lstm.py

@@ -1,26 +1,42 @@
 import logging
 
 from torch import nn
-from torch.nn.utils.rnn import pack_padded_sequence
+from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
 
 from chebai.models.base import ChebaiBaseNet
 
 logging.getLogger("pysmiles").setLevel(logging.CRITICAL)
 
 
 class ChemLSTM(ChebaiBaseNet):
-    def __init__(self, out_d, in_d, num_classes, criterion: nn.Module = None, num_layers=6, dropout=0.2, **kwargs):
+    def __init__(
+        self,
+        out_d,
+        in_d,
+        num_classes,
+        criterion: nn.Module = None,
+        num_layers=6,
+        dropout=0.2,
+        **kwargs,
+    ):
         super().__init__(
             out_dim=out_d,
             input_dim=in_d,
             criterion=criterion,
             num_classes=num_classes,
             **kwargs,
         )
-        self.lstm = nn.LSTM(in_d, out_d, batch_first=True, dropout=dropout, bidirectional=True, num_layers=num_layers)
+        self.lstm = nn.LSTM(
+            in_d,
+            out_d,
+            batch_first=True,
+            dropout=dropout,
+            bidirectional=True,
+            num_layers=num_layers,
+        )
         self.embedding = nn.Embedding(1400, in_d)
         self.output = nn.Sequential(
-            nn.Linear(out_d, out_d),
+            nn.Linear(out_d * 2, out_d),
             nn.ReLU(),
             nn.Dropout(0.2),
             nn.Linear(out_d, num_classes),
@@ -31,7 +47,9 @@ def forward(self, data, *args, **kwargs):
         x_lens = data["model_kwargs"]["lens"]
         x = self.embedding(x)
         x = pack_padded_sequence(x, x_lens, batch_first=True, enforce_sorted=False)
-        x = self.lstm(x)[1][0]
-        # = pad_packed_sequence(x, batch_first=True)[0]
+        x = self.lstm(x)[0]
+        x = pad_packed_sequence(x, batch_first=True)[0][
+            :, 0
+        ]  # reduce sequence dimension to first element
         x = self.output(x)
-        return x.squeeze(0)
+        return x