Added LSTM support

2025-03-26 02:43:05 -07:00
parent e6abdaebdc
commit fce22ba1dd
14 changed files with 721 additions and 5 deletions
@@ -1,7 +1,7 @@
 ## Update (3/25/2025) 

-GRU support has been added.
-Check out new example workflow at `examples/workflows/gru_classification.json` which achieves approximately **87% validation accuracy** for IMDB sentiment classification (which is in the expected range for GRU).
+LSTM & GRU support has been added.
+Check out new example workflows at `examples/workflows/lstm_classification.json` & `examples/workflows/gru_classification.json`. Both of them achieve approximately **87% validation accuracy** for IMDB sentiment classification (which is in the expected range for LSTM & GRU).

 ---
 # ComfyUI-Pt-Wrapper  
@@ -165,6 +165,8 @@ from .modules.pytorch_wrapper.ptn_l1_loss import PtnL1Loss
 from .modules.pytorch_wrapper.ptn_layer_norm import PtnLayerNorm
 from .modules.pytorch_wrapper.ptn_linear import PtnLinear
 from .modules.pytorch_wrapper.ptn_linear_model import PtnLinearModel
+from .modules.pytorch_wrapper.ptn_lstm import PtnLSTM
+from .modules.pytorch_wrapper.ptn_lstm_linear import PtnLSTMLinear
 from .modules.pytorch_wrapper.ptn_max_pool2d import PtnMaxPool2d
 from .modules.pytorch_wrapper.ptn_mse_loss import PtnMSELoss
 from .modules.pytorch_wrapper.ptn_nll_loss import PtnNLLLoss
@@ -361,6 +363,8 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
    "PtnInstanceNorm2d": PtnInstanceNorm2d,
    "PtnKLDivLoss": PtnKLDivLoss,
    "PtnL1Loss": PtnL1Loss,
+    "PtnLSTM": PtnLSTM,
+    "PtnLSTMLinear": PtnLSTMLinear,
    "PtnLayerNorm": PtnLayerNorm,
    "PtnLinear": PtnLinear,
    "PtnLinearModel": PtnLinearModel,
@@ -560,6 +564,8 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
    "PtnInstanceNorm2d": "Ptn Instance Norm 2d",
    "PtnKLDivLoss": "Ptn KL Div Loss",
    "PtnL1Loss": "Ptn L1 Loss",
+    "PtnLSTM": "Ptn LSTM",
+    "PtnLSTMLinear": "Ptn LSTM Linear",
    "PtnLayerNorm": "Ptn Layer Norm",
    "PtnLinear": "Ptn Linear",
    "PtnLinearModel": "Ptn Linear Model",
@@ -80,6 +80,7 @@ def format_display_name(class_name):
        (" N A", " NA"),
        ("R N N", "RNN"),
        ("G R U", "GRU"),
+        ("L S T M", "LSTM"),
        ("Adam W", "AdamW"),
        ("Re L U", "ReLU"),
        ("Si L U", "SiLU"),
@@ -158,6 +158,8 @@
 | [Ptn GRU](ptn_gru.md) | A gated recurrent unit (GRU) model consisting of one or more of a recurrent layer.   |
 | [Ptn GRU Linear](ptn_gru_linear.md) | A recurrent neural network (GRU) model with a linear head. |
 | [Ptn Instance Norm 2d](ptn_instance_norm2d.md) | A normalization model to normalize elements over spatial axes within each channel for each sample. |
+| [Ptn LSTM](ptn_lstm.md) | A long short-term memory  (LSTM) model consisting of one or more of a recurrent layer.   |
+| [Ptn LSTM Linear](ptn_lstm_linear.md) | A recurrent neural network (LSTM) model with a linear head. |
 | [Ptn Layer Norm](ptn_layer_norm.md) | A normalization model to normalize elements over specified axes. |
 | [Ptn Linear](ptn_linear.md) | A linear model consisting of a single dense layer.   |
 | [Ptn Linear Model](ptn_linear_model.md) | A linear model consisting of dense layers.   |
@@ -0,0 +1,34 @@
+# Ptn LSTM
+A long short-term memory  (LSTM) model consisting of one or more of a recurrent layer.  
+
+    Args:
+        input_size (int): The number of input features.  
+        hidden_size (int): The number of output features of the hidden layer matrix.  
+        num_layers (int): Number of hidden layers.  
+        bias (bool): Use bias or not.  
+        batch_first (bool): Indicates the shape of the input. Input is assumed to be a rank-3 tensor. If True, the shape is [Batch, Seq, Token]. If False, the shape is [Seq, Batch, Token]. Note that default is set to `True` unlike PyTorch's LSTM default model parameter.
+        dropout (float): Dropout probability to be applied to the output of intermediate layers if non-zero. This will not be applied to the last layer's output.  
+        bidirectional (bool): Processes the input in a backward direction and append the result to the output of the forward direction.  
+        proj_size (int): If set to greater than 0, applies a projection to each output hidden state, reducing its dimensionality to proj_size.
+
+## Input
+| Name | Data type |
+|---|---|
+| input_size | Int |
+| hidden_size | Int |
+| num_layers | Int |
+| bias | Boolean |
+| batch_first | Boolean |
+| dropout | Float |
+| bidirectional | Boolean |
+| proj_size | Int |
+
+## Output
+| Data type |
+|---|
+| Ptmodel |
+
+<HR>
+Category: PyTorch wrapper - Model
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
@@ -0,0 +1,38 @@
+# Ptn LSTM Linear
+A recurrent neural network (LSTM) model with a linear head.
+
+    Args:
+        input_size (int): The number of input features.
+        hidden_size (int): The number of output features of the hidden layer matrix.
+        num_layers (int): Number of hidden layers in LSTM.
+        bias (bool): Use bias or not.
+        batch_first (bool): Indicates the shape of the input. Input is assumed to be a rank-3 tensor. If True, the shape is [Batch, Seq, Token]. If False, the shape is [Seq, Batch, Token]. Note that default is set to `True` unlike PyTorch's LSTM default model parameter.
+        dropout (float): Dropout probability to be applied to the output of intermediate layers if non-zero. This will not be applied to the last layer's output.
+        bidirectional (bool): Processes the input in a backward direction and append the result to the output of the forward direction.
+        proj_size (int): If set to greater than 0, applies a projection to each output hidden state, reducing its dimensionality to proj_size.  
+        linear_output_size (int): The number of output features of the linear layer.
+        linar_bias (bool): Use bias or not in linear layer.
+
+## Input
+| Name | Data type |
+|---|---|
+| input_size | Int |
+| hidden_size | Int |
+| num_layers | Int |
+| bias | Boolean |
+| batch_first | Boolean |
+| dropout | Float |
+| bidirectional | Boolean |
+| proj_size | Int |
+| linear_output_size | Int |
+| linear_bias | Boolean |
+
+## Output
+| Data type |
+|---|
+| Ptmodel |
+
+<HR>
+Category: PyTorch wrapper - Model
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
@@ -3,7 +3,7 @@
 You can use **ComfyUI-Pt-Wrapper** to train an RNN for text classification, such as sentiment analysis on a movie review dataset.

 You can find the workflow at `examples/workflows/rnn_classification.json`.
-You can also use GRU workflow at `examples/workflows/gru_classification.json`.
+You can also use GRU workflow at `examples/workflows/gru_classification.json` or LSTM workflow at `examples/workflows/lstm_classification.json`.

 ![Workflow](images/rnn_classification.png)

@@ -43,4 +43,4 @@ You can also enable `classification_metrics` to print validation accuracy to the
 ## Expected Accuracy

 Using this workflow, you should get around **83% validation accuracy** on the Hugging Face version of the IMDB dataset, which is a solid result for an RNN.
-The GRU workflow achieves approximately **87% validation accuracy**.
+Both GRU & LSTM workflows achieve approximately **87% validation accuracy**.
@@ -0,0 +1,91 @@
+from typing import Any, Dict, Tuple
+import torch
+import torch.nn as nn
+
+
+class PtnLSTM:
+    """
+    Ptn LSTM:
+    A long short-term memory  (LSTM) model consisting of one or more of a recurrent layer.  
+
+        Args:
+            input_size (int): The number of input features.  
+            hidden_size (int): The number of output features of the hidden layer matrix.  
+            num_layers (int): Number of hidden layers.  
+            bias (bool): Use bias or not.  
+            batch_first (bool): Indicates the shape of the input. Input is assumed to be a rank-3 tensor. If True, the shape is [Batch, Seq, Token]. If False, the shape is [Seq, Batch, Token]. Note that default is set to `True` unlike PyTorch's LSTM default model parameter.
+            dropout (float): Dropout probability to be applied to the output of intermediate layers if non-zero. This will not be applied to the last layer's output.  
+            bidirectional (bool): Processes the input in a backward direction and append the result to the output of the forward direction.  
+            proj_size (int): If set to greater than 0, applies a projection to each output hidden state, reducing its dimensionality to proj_size.  
+
+    category: PyTorch wrapper - Model
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required input types.
+        """
+        return {
+            "required": {
+                "input_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
+                "hidden_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
+                "num_layers": ("INT", {"default": 1, "min": 1, "max": 1e3}),
+                "bias": ("BOOLEAN", {"default": True}),
+                "batch_first": ("BOOLEAN", {"default": True}),
+                "dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1}),
+                "bidirectional": ("BOOLEAN", {"default": False}),
+                "proj_size": ("INT", {"default": 0, "min": 0, "max": 1e6}),
+            }
+        }
+
+    RETURN_TYPES: Tuple[str] = ("PTMODEL",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self,
+          input_size: int,
+          hidden_size: int,
+          num_layers: int,
+          bias: bool,
+          batch_first: bool,
+          dropout: float,
+          bidirectional: bool,
+          proj_size: int
+        ) -> Tuple[nn.Module]:
+        """
+        Constructs an LSTM model.
+
+        Args:
+            input_size (int): The number of input features.
+            hidden_size (int): The number of output features of the hidden layer matrix.
+            num_layers (int): Number of hidden layers.
+            bias (bool): Use bias or not.
+            batch_first (bool): Indicates the shape of the input. Input is assumed to be a rank-3 tensor. If True, the shape is [Batch, Seq, Token]. If False, the shape is [Seq, Batch, Token].
+            dropout (float): Dropout probability to be applied to the output of intermediate layers if non-zero. This will not be applied to the last layer's output.
+            bidirectional (bool): Processes the input in a backward direction and append the result to the output of the forward direction.
+            proj_size (int): If set to greater than 0, applies a projection to each output hidden state, reducing its dimensionality to proj_size.
+
+         Returns:
+            Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
+        """
+        with torch.inference_mode(False):
+            model = nn.LSTM(
+                input_size=input_size,
+                hidden_size=hidden_size,
+                num_layers=num_layers,
+                bias=bias,
+                batch_first=batch_first,
+                dropout=dropout,
+                bidirectional=bidirectional,
+                proj_size=proj_size
+            )
+
+        return (model,)
@@ -0,0 +1,170 @@
+from typing import Any, Dict, Tuple
+import torch
+import torch.nn as nn
+from pytorch_wrapper.rnn_output_reshape import extract_rnn_return_value_with_adjusted_label
+
+
+class LSTMLinear(nn.Module):
+    """
+    A long short-term memory  (LSTM) model with a linear output layer.
+    
+    This module uses a standard PyTorch LSTM, followed by a linear layer applied to the final output tensor.
+
+    pragma: skip_doc
+    """
+
+    def __init__(self,
+        input_size: int,
+        hidden_size: int,
+        num_layers: int,
+        bias: bool,
+        batch_first: bool,
+        dropout: float,
+        bidirectional: bool,
+        proj_size: int,
+        linear_output_size: int,
+        linear_bias: bool
+    ):
+        super().__init__()
+        self.lstm_model = nn.LSTM(
+            input_size=input_size,
+            hidden_size=hidden_size,
+            num_layers=num_layers,
+            bias=bias,
+            batch_first=batch_first,
+            dropout=dropout,
+            bidirectional=bidirectional,
+            proj_size=proj_size
+        )
+        self.bidirectional = bidirectional
+        self.batch_first = batch_first
+        directions = 2 if bidirectional else 1
+
+        self.linear_model = nn.Linear(
+            in_features=hidden_size * directions,
+            out_features=linear_output_size,
+            bias=linear_bias
+        )
+        torch.nn.init.xavier_uniform_(self.linear_model.weight)
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass through the LSTM and the linear layer.
+
+        Args:
+            inputs (torch.Tensor): Input tensor of shape (batch, seq, input_size) if batch_first=True.
+
+        Returns:
+            torch.Tensor: Output of the linear layer.
+        """
+        out_seq, _ = self.lstm_model(inputs)
+
+        x = extract_rnn_return_value_with_adjusted_label(
+            inputs,  # This is needed to find out padded tokens
+            None,
+            out_seq,
+            self.bidirectional,
+            self.batch_first,
+            False,
+            return_valid_token_mean=True
+        )
+
+        outputs = self.linear_model(x)
+        return outputs
+
+
+class PtnLSTMLinear:
+    """
+    Ptn LSTM Linear:
+    A recurrent neural network (LSTM) model with a linear head.
+    
+        Args:
+            input_size (int): The number of input features.
+            hidden_size (int): The number of output features of the hidden layer matrix.
+            num_layers (int): Number of hidden layers in LSTM.
+            bias (bool): Use bias or not.
+            batch_first (bool): Indicates the shape of the input. Input is assumed to be a rank-3 tensor. If True, the shape is [Batch, Seq, Token]. If False, the shape is [Seq, Batch, Token]. Note that default is set to `True` unlike PyTorch's LSTM default model parameter.
+            dropout (float): Dropout probability to be applied to the output of intermediate layers if non-zero. This will not be applied to the last layer's output.
+            bidirectional (bool): Processes the input in a backward direction and append the result to the output of the forward direction.
+            proj_size (int): If set to greater than 0, applies a projection to each output hidden state, reducing its dimensionality to proj_size.  
+            linear_output_size (int): The number of output features of the linear layer.
+            linar_bias (bool): Use bias or not in linear layer.
+
+    category: PyTorch wrapper - Model
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required input types.
+        """
+        return {
+            "required": {
+                "input_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
+                "hidden_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
+                "num_layers": ("INT", {"default": 1, "min": 1, "max": 1e3}),
+                "bias": ("BOOLEAN", {"default": True}),
+                "batch_first": ("BOOLEAN", {"default": True}),
+                "dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1}),
+                "bidirectional": ("BOOLEAN", {"default": False}),
+                "proj_size": ("INT", {"default": 0, "min": 0, "max": 1e6}),
+                "linear_output_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
+                "linear_bias": ("BOOLEAN", {"default": True}),
+            }
+        }
+
+    RETURN_TYPES: Tuple[str] = ("PTMODEL",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self,
+          input_size: int,
+          hidden_size: int,
+          num_layers: int,
+          bias: bool,
+          batch_first: bool,
+          dropout: float,
+          bidirectional: bool,
+          proj_size: int,
+          linear_output_size: int,
+          linear_bias: bool
+        ) -> Tuple[nn.Module]:
+        """
+        Constructs an LSTM model.
+
+        Args:
+            input_size (int): The number of input features.
+            hidden_size (int): The number of output features of the hidden layer matrix.
+            num_layers (int): Number of hidden layers.
+            bias (bool): Use bias or not.
+            batch_first (bool): Indicates the shape of the input. Input is assumed to be a rank-3 tensor. If True, the shape is [Batch, Seq, Token]. If False, the shape is [Seq, Batch, Token].
+            dropout (float): Dropout probability to be applied to the output of intermediate layers if non-zero. This will not be applied to the last layer's output.
+            bidirectional (bool): Processes the input in a backward direction and append the result to the output of the forward direction.
+            proj_size (int): If set to greater than 0, applies a projection to each output hidden state, reducing its dimensionality to proj_size.  
+            linear_output_size (int): The number of output features of the linear layer.
+            linar_bias (bool): Use bias or not in linear layer.
+
+        Returns:
+            Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
+        """
+        with torch.inference_mode(False):
+            model = LSTMLinear(
+                input_size=input_size,
+                hidden_size=hidden_size,
+                num_layers=num_layers,
+                bias=bias,
+                batch_first=batch_first,
+                dropout=dropout,
+                bidirectional=bidirectional,
+                proj_size=proj_size,
+                linear_output_size=linear_output_size,
+                linear_bias=linear_bias
+            )
+            return (model,)
@@ -1,7 +1,7 @@
 [project]
 name = "pt-wrapper"
 description = "PyTorch extension for ComfyUI featuring extensive PyTorch wrapper nodes for seamless tensor operations and PyTorch model training."
-version = "3.3.1"
+version = "3.3.2"
 license = {file = "LICENSE"}

 [project.urls]
@@ -0,0 +1,126 @@
+"""
+Dataset URL:https://huggingface.co/datasets/stanfordnlp/imdb
+"""
+import os
+import sys
+import unittest
+import torch
+import torch.nn as nn
+import random
+import numpy as np
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+
+from pytorch_wrapper.ptn_lstm_linear import PtnLSTMLinear
+from pytorch_wrapper.pto_adamw import PtoAdamW
+from pytorch_wrapper.ptv_hf_glove_dataset import PtvHfGloveDataset
+from pytorch_wrapper.pt_data_loader import PtDataLoader
+from pytorch_wrapper.pt_train_rnn_model import PtTrainRNNModel
+from pytorch_wrapper.pt_save_model import PtSaveModel
+from pytorch_wrapper.pt_load_model import PtLoadModel
+from pytorch_wrapper.pto_lr_scheduler_reduce_on_plateau import PtoLrSchedulerReduceOnPlateau
+from pytorch_wrapper.ptn_bce_with_logits_loss import PtnBCEWithLogitsLoss
+from pytorch_wrapper.utils import set_seed
+
+
+class TestPtTrainLSTMLinearTextClassificationModel(unittest.TestCase):
+
+    @unittest.skipIf(os.getenv("RUN_SKIPPED_TESTS") != "1", "Skipping unless explicitly enabled")
+    def setUp(self):
+        set_seed(42)
+
+        self.loss_fun =PtnBCEWithLogitsLoss().f(reduction="mean")[0]
+        dataset_node = PtvHfGloveDataset()
+        self.embedding_dim = 100
+        self.max_seq_len = 300
+        train_dataset = dataset_node.f(
+            "imdb",  # dataset_name
+            "train",  #split
+            self.embedding_dim,
+            self.max_seq_len,
+            "text",  # sample field name in the dataset
+            "label"  # label field name in the dataset
+        )[0]
+        test_dataset = dataset_node.f(
+            "imdb",  # dataset_name
+            "test",  #split
+            self.embedding_dim,
+            self.max_seq_len,
+            "text",  # sample field name in the dataset
+            "label"  # label field name in the dataset
+        )[0]
+        train_data_loader_node = PtDataLoader()
+        test_data_loader_node = PtDataLoader()
+
+        self.train_loader = train_data_loader_node.f(
+            dataset=train_dataset,
+            batch_size=128,
+            shuffle=True,
+            parameters='{"num_workers":1}'
+        )[0]
+
+        self.test_loader = test_data_loader_node.f(
+            dataset=test_dataset,
+            batch_size=128,
+            shuffle=False,
+            parameters='{"num_workers":1}'
+        )[0]
+
+        self.model_path = "lstm_text_classification_test.pt"
+
+        self.model_node = PtnLSTMLinear()
+        self.model = self.model_node.f(
+            self.embedding_dim, # input_size
+            512, # hidden_size
+            2, # num_layers
+            True, # bias,
+            True, # batch_first [B, Seq, Token]
+            0.4, # dropout
+            True, # bidirectional,
+            0,  # proj_size
+            1, # linear output size
+            True, # linear bias
+        )[0]
+
+        self.optimizer = PtoAdamW().f(self.model,
+                                      0.001,
+                                      0.9, 0.999,
+                                      weight_decay=0.01,
+                                      amsgrad=False)[0]
+        
+        self.scheduler = PtoLrSchedulerReduceOnPlateau().f(self.optimizer,
+                                                grace_period=10,
+                                                gamma=0.5
+                                               )[0]
+
+        self.trainer = PtTrainRNNModel()
+        self.save_model = PtSaveModel()
+        self.load_model = PtLoadModel()
+        self.linear_head = True
+        self.epochs = 50
+
+    @unittest.skipIf(os.getenv("RUN_SKIPPED_TESTS") != "1", "Skipping unless explicitly enabled")
+    def test_1(self):
+
+        trained_model = self.trainer.f(
+                        self.model,
+                        self.linear_head,
+                        self.train_loader,
+                        self.optimizer,
+                        self.loss_fun,
+                        self.epochs, # epochs
+                        use_gpu=True,
+                        early_stopping=False,
+                        early_stopping_rounds=1,
+                        scheduler=self.scheduler,
+                        output_best_val_model=False,
+                        classification_metrics=True,
+                        use_valid_token_mean=True,
+                        val_loader=self.test_loader)[0]
+        self.save_model.f(trained_model, self.model_path)
+
+
+if __name__ == "__main__":
+    unittest.main()
@@ -0,0 +1,91 @@
+import os
+import sys
+import unittest
+import torch
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+from pytorch_wrapper.ptn_lstm_linear import PtnLSTMLinear
+
+
+class TestPtnLSTMLinear(unittest.TestCase):
+    def setUp(self):
+        """Set up test instance."""
+        self.model_node = PtnLSTMLinear()
+
+    def test_seq_first_1_layer_uni(self):
+        self.model = self.model_node.f(
+            2, # input_size
+            3, # hidden_size
+            1, # num_layers
+            True, # bias,
+            False, # batch_first so [Seq, B, Token]
+            0.0, # dropout
+            False, # bidirectional,
+            0,  # proj_size
+            5, # linear output features
+            True # linear bias
+        )[0]
+
+        x = torch.randn(10, 8, 2)  # [Seq, B, Token]
+        outputs= self.model(x)
+        self.assertEqual(outputs.size(), torch.Size([8, 5]))
+
+    def test_seq_first_1_layer_bidi(self):
+        self.model = self.model_node.f(
+            2, # input_size
+            3, # hidden_size
+            1, # num_layers
+            True, # bias,
+            False, # batch_first so [Seq, B, Token]
+            0.0, # dropout
+            True, # bidirectional,
+            0,  # proj_size
+            5, # linear output features
+            True # linear bias
+        )[0]
+
+        x = torch.randn(10, 8, 2)  # [Seq, B, Token]
+        outputs= self.model(x)
+        self.assertEqual(outputs.size(), torch.Size([8, 5]))
+
+    def test_seq_first_2_layer_uni(self):
+        self.model = self.model_node.f(
+            2, # input_size
+            3, # hidden_size
+            2, # num_layers
+            True, # bias,
+            False, # batch_first so [Seq, B, Token]
+            0.0, # dropout
+            False, # bidirectional,
+            0,  # proj_size
+            5, # linear output features
+            True # linear bias
+        )[0]
+
+        x = torch.randn(10, 8, 2)  # [Seq, B, Token]
+        outputs= self.model(x)
+        self.assertEqual(outputs.size(), torch.Size([8, 5]))
+
+    def test_batch_first_1_layer_uni(self):
+        self.model = self.model_node.f(
+            2, # input_size
+            3, # hidden_size
+            1, # num_layers
+            True, # bias,
+            True, # batch_first
+            0.0, # dropout
+            False, # bidirectional,
+            0,  # proj_size
+            5, # linear output features
+            True # linear bias
+        )[0]
+
+        x = torch.randn(8, 10, 2)  # [B, Seq, Token]
+        outputs= self.model(x)
+        self.assertEqual(outputs.size(), torch.Size([8, 5]))
+
+
+if __name__ == "__main__":
+    unittest.main()
@@ -0,0 +1,156 @@
+import os
+import sys
+import unittest
+import torch
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+from pytorch_wrapper.ptn_lstm import PtnLSTM
+
+
+class TestPtnLSTM(unittest.TestCase):
+    def setUp(self):
+        """Set up test instance."""
+        self.model_node = PtnLSTM()
+
+    def test_shape_1_layer(self):
+        self.model = self.model_node.f(
+            2, # input_size
+            3, # hidden_size
+            1, # num_layers
+            True, # bias,
+            False, # batch_first so [Seq, B, Token]
+            0.0, # dropout
+            False, # bidirectional,
+            0  # proj_size
+        )[0]
+
+        x = torch.randn(10, 8, 2)  # [Seq, B, Token]
+        out_seq_last_layer, (out_last_value, cn) = self.model(x)
+        self.assertEqual(out_last_value.size(), torch.Size([1, 8, 3]))  # num_layers, batch, hidden_size
+        self.assertEqual(out_seq_last_layer.size(), torch.Size([10, 8, 3]))  # seq, batch, hidden_size
+
+    def test_shape_2_layers(self):
+        self.model = self.model_node.f(
+            2, # input_size
+            3, # hidden_size
+            2, # num_layers
+            True, # bias,
+            False, # batch_first so [Seq, B, Token]
+            0.0, # dropout
+            False, # bidirectional
+            0  # proj_size
+        )[0]
+
+        x = torch.randn(10, 8, 2)  # [Seq, B, Token]
+        out_seq_last_layer, (out_last_value, cn) = self.model(x)
+        self.assertEqual(out_last_value.size(), torch.Size([2, 8, 3]))  # num_layers, batch, hidden_size
+        self.assertEqual(out_seq_last_layer.size(), torch.Size([10, 8, 3]))  # seq, batch, hidden_size
+
+    def test_bidirectional_1_layer(self):
+        self.model = self.model_node.f(
+            2, # input_size
+            3, # hidden_size
+            1, # num_layers
+            True, # bias,
+            False, # batch_first so [Seq, B, Token]
+            0.0, # dropout
+            True, # bidirectional
+            0  # proj_size
+        )[0]
+
+        x = torch.randn(10, 8, 2)  # [Seq, B, Token]
+        out_seq_last_layer, (out_last_value, cn) = self.model(x)
+        self.assertEqual(out_last_value.size(), torch.Size([2, 8, 3]))  # 2 * num_layers, batch, hidden_size
+        self.assertEqual(out_seq_last_layer.size(), torch.Size([10, 8, 6]))  # seq, batch, 2 * hidden_size
+
+    def test_bidirectional_2_layers(self):
+        self.model = self.model_node.f(
+            2, # input_size
+            3, # hidden_size
+            2, # num_layers
+            True, # bias,
+            False, # batch_first so [Seq, B, Token]
+            0.0, # dropout
+            True, # bidirectional
+            0  # proj_size
+        )[0]
+
+        x = torch.randn(10, 8, 2)  # [Seq, B, Token]
+        out_seq_last_layer, (out_last_value, cn) = self.model(x)
+        self.assertEqual(out_last_value.size(), torch.Size([4, 8, 3]))  # 2 * num_layers, batch, hidden_size
+        self.assertEqual(out_seq_last_layer.size(), torch.Size([10, 8, 6]))  # seq, batch, 2 * hidden_size
+
+    # batch-first "bf" test cases
+    def test_bf_shape_1_layer(self):
+        self.model = self.model_node.f(
+            2, # input_size
+            3, # hidden_size
+            1, # num_layers
+            True, # bias,
+            True, # batch_first so [B, Seq, Token]
+            0.0, # dropout
+            False, # bidirectional
+            0  # proj_size
+        )[0]
+
+        x = torch.randn(8, 10, 2)  # [B, Seq, Token]
+        out_seq_last_layer, (out_last_value, cn) = self.model(x)
+        self.assertEqual(out_last_value.size(), torch.Size([1, 8, 3]))  # num_layers, batch, hidden_size
+        self.assertEqual(out_seq_last_layer.size(), torch.Size([8, 10, 3]))  # seq, batch, hidden_size
+
+    def test_bf_shape_2_layers(self):
+        self.model = self.model_node.f(
+            2, # input_size
+            3, # hidden_size
+            2, # num_layers
+            True, # bias,
+            True, # batch_first so [B, Seq, Token]
+            0.0, # dropout
+            False, # bidirectional
+            0  # proj_size
+        )[0]
+
+        x = torch.randn(8, 10, 2)  # [B, Seq, Token]
+        out_seq_last_layer, (out_last_value, cn) = self.model(x)
+        self.assertEqual(out_last_value.size(), torch.Size([2, 8, 3]))  # num_layers, batch, hidden_size
+        self.assertEqual(out_seq_last_layer.size(), torch.Size([8, 10, 3]))  # batch, seq, hidden_size
+
+    def test_bf_bidirectional_1_layer(self):
+        self.model = self.model_node.f(
+            2, # input_size
+            3, # hidden_size
+            1, # num_layers
+            True, # bias,
+            True, # batch_first so [B, Seq, Token]
+            0.0, # dropout
+            True, # bidirectional
+            0  # proj_size
+        )[0]
+
+        x = torch.randn(8, 10, 2)  # [B, Seq, Token]
+        out_seq_last_layer, (out_last_value, cn) = self.model(x)
+        self.assertEqual(out_last_value.size(), torch.Size([2, 8, 3]))  # 2 * num_layers, batch, hidden_size
+        self.assertEqual(out_seq_last_layer.size(), torch.Size([8, 10, 6]))  # batch, seq, 2 * hidden_size
+
+    def test_bf_bidirectional_2_layers(self):
+        self.model = self.model_node.f(
+            2, # input_size
+            3, # hidden_size
+            2, # num_layers
+            True, # bias,
+            True, # batch_first so [B, Seq, Token]
+            0.0, # dropout
+            True, # bidirectional
+            0  # proj_size
+        )[0]
+
+        x = torch.randn(8, 10, 2)  # [B, Seq, Token]
+        out_seq_last_layer, (out_last_value, cn) = self.model(x)
+        self.assertEqual(out_last_value.size(), torch.Size([4, 8, 3]))  # 2 * num_layers, batch, hidden_size
+        self.assertEqual(out_seq_last_layer.size(), torch.Size([8, 10, 6]))  # batch, seq, 2 * hidden_size
+
+
+if __name__ == "__main__":
+    unittest.main()