Added GRU support.

README.md

@@ -1,3 +1,9 @@
+## 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).
+
+---
 # ComfyUI-Pt-Wrapper  
 
 ComfyUI-Pt-Wrapper is a spin-off of [ComfyUI-Data-Analysis](https://github.com/HowToSD/ComfyUI-Data-Analysis), designed as a dedicated extension for PyTorch wrapper nodes originally included in the analysis extension.

__init__.py

@@ -156,6 +156,8 @@ from .modules.pytorch_wrapper.ptn_chained_model import PtnChainedModel
 from .modules.pytorch_wrapper.ptn_conv2d import PtnConv2d
 from .modules.pytorch_wrapper.ptn_conv_model import PtnConvModel
 from .modules.pytorch_wrapper.ptn_cross_entropy_loss import PtnCrossEntropyLoss
+from .modules.pytorch_wrapper.ptn_gru import PtnGRU
+from .modules.pytorch_wrapper.ptn_gru_linear import PtnGRULinear
 from .modules.pytorch_wrapper.ptn_huber_loss import PtnHuberLoss
 from .modules.pytorch_wrapper.ptn_instance_norm2d import PtnInstanceNorm2d
 from .modules.pytorch_wrapper.ptn_kl_div_loss import PtnKLDivLoss
@@ -353,6 +355,8 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
     "PtnConv2d": PtnConv2d,
     "PtnConvModel": PtnConvModel,
     "PtnCrossEntropyLoss": PtnCrossEntropyLoss,
+    "PtnGRU": PtnGRU,
+    "PtnGRULinear": PtnGRULinear,
     "PtnHuberLoss": PtnHuberLoss,
     "PtnInstanceNorm2d": PtnInstanceNorm2d,
     "PtnKLDivLoss": PtnKLDivLoss,
@@ -550,6 +554,8 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
     "PtnConv2d": "Ptn Conv 2d",
     "PtnConvModel": "Ptn Conv Model",
     "PtnCrossEntropyLoss": "Ptn Cross Entropy Loss",
+    "PtnGRU": "Ptn GRU",
+    "PtnGRULinear": "Ptn GRU Linear",
     "PtnHuberLoss": "Ptn Huber Loss",
     "PtnInstanceNorm2d": "Ptn Instance Norm 2d",
     "PtnKLDivLoss": "Ptn KL Div Loss",

dev/scripts/init_py_generate.py

@@ -79,6 +79,7 @@ def format_display_name(class_name):
         ("SN S", "SNS"),
         (" N A", " NA"),
         ("R N N", "RNN"),
+        ("G R U", "GRU"),
         ("Adam W", "AdamW"),
         ("Re L U", "ReLU"),
         ("Si L U", "SiLU"),

docs/reference/node_reference.md

@@ -155,6 +155,8 @@
 | [Ptn Chained Model](ptn_chained_model.md) | Constructs a chained PyTorch model. |
 | [Ptn Conv 2d](ptn_conv2d.md) |     A convolutional model consisting of a single conv2d layer.   |
 | [Ptn Conv Model](ptn_conv_model.md) | A convolutional model consisting of multiple convolutional layers.   |
+| [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 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.   |
@@ -228,7 +230,7 @@
 | [Pt Train Classification Model](pt_train_classification_model.md) | Trains a classification model using a given dataset, optimizer, and number of epochs.   |
 | [Pt Train Classification Model Lr](pt_train_classification_model_lr.md) | Trains a classification model using a given dataset, optimizer, and number of epochs with learning rate decay.   |
 | [Pt Train Model](pt_train_model.md) | Trains a model using a given dataset, loss function, optimizer, and number of epochs with learning rate decay.   |
-| [Pt Train R N N Model](pt_train_rnn_model.md) | Trains an RNN model using a given dataset, loss function, optimizer, and number of epochs with learning rate decay.   |
+| [Pt Train RNN Model](pt_train_rnn_model.md) | Trains an RNN model using a given dataset, loss function, optimizer, and number of epochs with learning rate decay.   |
 | [Pt Train Regression Model](pt_train_regression_model.md) | Trains a regression model using a given dataset, optimizer, and number of epochs.   |
 | [Pto Lr Scheduler Cosine Annealing](pto_lr_scheduler_cosine_annealing.md) | Creates a cosine annealing learning rate scheduler for an optimizer. |
 | [Pto Lr Scheduler Reduce On Plateau](pto_lr_scheduler_reduce_on_plateau.md) | Creates a reduce-on-plateau learning rate scheduler for an optimizer. |

docs/reference/pt_train_rnn_model.md

@@ -1,4 +1,4 @@
-# Pt Train R N N Model
+# Pt Train RNN Model
 Trains an RNN model using a given dataset, loss function, optimizer, and number of epochs with learning rate decay.  
 
    Args:  
@@ -15,7 +15,8 @@ Trains an RNN model using a given dataset, loss function, optimizer, and number
         classification_metrics (bool): If True and if val_loader is set, print out classification metrics.  This is only valid for classification model training.  
         use_valid_token_mean (bool): Computes the mean of the RNN output for non-zero inputs and feeds to the linear layer.
         val_loader (PTDATALOADER): DataLoader for validation dataset (optional).  
-        scheduler (PTLRSCHEDULER): Learning rate scheduler.
+        scheduler (PTLRSCHEDULER): Learning rate scheduler.  
+        h_0 (TENSOR): Optional initial hidden state. For unidirectional models, the shape is [num_layers, batch_size, hidden_size]. For bidirectional models, it is [2 * num_layers, batch_size, hidden_size]. If not provided, a zero tensor is used.
 
 ## Input
 | Name | Data type |
@@ -34,6 +35,7 @@ Trains an RNN model using a given dataset, loss function, optimizer, and number
 | use_valid_token_mean | Boolean |
 | scheduler | Ptlrscheduler |
 | val_loader | Ptdataloader |
+| h_0 | Tensor |
 
 ## Output
 | Data type |

docs/reference/ptn_gru.md

@@ -0,0 +1,32 @@
+# Ptn GRU
+A gated recurrent unit (GRU) 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 GRU 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.
+
+## Input
+| Name | Data type |
+|---|---|
+| input_size | Int |
+| hidden_size | Int |
+| num_layers | Int |
+| bias | Boolean |
+| batch_first | Boolean |
+| dropout | Float |
+| bidirectional | Boolean |
+
+## Output
+| Data type |
+|---|
+| Ptmodel |
+
+<HR>
+Category: PyTorch wrapper - Model
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

docs/reference/ptn_gru_linear.md

@@ -0,0 +1,36 @@
+# Ptn GRU Linear
+A recurrent neural network (GRU) 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 GRU.
+        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 GRU 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.
+        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 |
+| 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.

docs/reference/ptn_rnn_linear.md

@@ -24,6 +24,8 @@ A recurrent neural network (RNN) model with a linear head.
 | batch_first | Boolean |
 | dropout | Float |
 | bidirectional | Boolean |
+| linear_output_size | Int |
+| linear_bias | Boolean |
 
 ## Output
 | Data type |

docs/training_rnn_for_classification.md

@@ -3,6 +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`.
 
 ![Workflow](images/rnn_classification.png)
 
@@ -42,3 +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**.

modules/pytorch_wrapper/hf_glove_dataset.py

@@ -72,6 +72,9 @@ class HfGloveDataset(Dataset):
         # This strips punctuation like "!", ".", "(", but keeps words and numbers
         text = re.sub(r"[^a-zA-Z0-9\s]", " ", text)
 
+        # Replace consecutive white space to a single white space
+        text = re.sub(r"  +", " ", text)
+
         # Convert to lowercase and split by whitespace to get token list
         return text.lower().split()
 

modules/pytorch_wrapper/pt_train_rnn_model.py

@@ -9,6 +9,7 @@ from .rnn_output_reshape import extract_rnn_return_value_with_adjusted_label
 
 class PtTrainRNNModel:
     """
+    Pt Train RNN Model:
     Trains an RNN model using a given dataset, loss function, optimizer, and number of epochs with learning rate decay.  
 
        Args:  
@@ -26,6 +27,7 @@ class PtTrainRNNModel:
             use_valid_token_mean (bool): Computes the mean of the RNN output for non-zero inputs and feeds to the linear layer.
             val_loader (PTDATALOADER): DataLoader for validation dataset (optional).  
             scheduler (PTLRSCHEDULER): Learning rate scheduler.  
+            h_0 (TENSOR): Optional initial hidden state. For unidirectional models, the shape is [num_layers, batch_size, hidden_size]. For bidirectional models, it is [2 * num_layers, batch_size, hidden_size]. If not provided, a zero tensor is used.
 
     category: PyTorch wrapper - Training
     """
@@ -56,6 +58,7 @@ class PtTrainRNNModel:
             "optional": {
                 "scheduler": ("PTLRSCHEDULER", {}),
                 "val_loader": ("PTDATALOADER", {}),
+                "h_0": ("TENSOR", {}),
             }
         }
 
@@ -82,7 +85,8 @@ class PtTrainRNNModel:
           classification_metrics:bool,
           use_valid_token_mean: bool,
           scheduler: Optional[Union[torch.optim.lr_scheduler._LRScheduler, torch.optim.lr_scheduler.ReduceLROnPlateau]] = None,
-          val_loader: Optional[torch.utils.data.DataLoader] = None) -> Tuple[nn.Module, torch.Tensor, torch.Tensor]:
+          val_loader: Optional[torch.utils.data.DataLoader] = None,
+          h_0: Optional[torch.Tensor] = None) -> Tuple[nn.Module, torch.Tensor, torch.Tensor]:
         """
         Trains an RNN model.
 
@@ -101,6 +105,7 @@ class PtTrainRNNModel:
             use_valid_token_mean (bool): Computes the mean of the RNN output for non-zero inputs and feeds to the linear layer.
             val_loader (Optional[torch.utils.data.DataLoader]): DataLoader for validation dataset (optional).
             scheduler (Optional[Union[torch.optim.lr_scheduler._LRScheduler, torch.optim.lr_scheduler.ReduceLROnPlateau]]): Learning rate scheduler.
+            h_0 (optional[torch.Tensor]): Optional initial hidden state. For unidirectional models, the shape is [num_layers, batch_size, hidden_size]. For bidirectional models, it is [2 * num_layers, batch_size, hidden_size]. If not provided, a zero tensor is used.
 
         Returns:
             Tuple[torch.nn.Module, torch.Tensor, torch.Tensor]: A tuple containing the trained model, 
@@ -142,14 +147,20 @@ class PtTrainRNNModel:
                         y = y.to("cuda")
                     optimizer.zero_grad()
                     if linear_head:
-                        y_hat = model(x)
+                        if h_0:
+                            y_hat = model(x, hx=h_0)  # Note that both RNN and GRU use hx. See PyTorch code. https://github.com/pytorch/pytorch/blob/v2.6.0/torch/nn/modules/rnn.py
+                        else:
+                            y_hat = model(x)
                         if y.dim() == 3:
                             y = y[:,-1,:]  # Grab last token
                         elif y.dim() == 1:
                             y = torch.unsqueeze(y, dim=-1)
                             y = y.to(torch.float32)
                     else:
-                        output_seq, _ = model(x)
+                        if h_0:
+                            output_seq, _ = model(x, hx=h_0)
+                        else:
+                            output_seq, _ = model(x)
                         outputs, adjusted_labels = extract_rnn_return_value_with_adjusted_label(
                             x,
                             y,

modules/pytorch_wrapper/ptn_gru.py

@@ -0,0 +1,88 @@
+from typing import Any, Dict, Tuple
+import torch
+import torch.nn as nn
+
+
+class PtnGRU:
+    """
+    Ptn GRU:
+    A gated recurrent unit (GRU) 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 GRU 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.
+
+    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}),
+            }
+        }
+
+    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
+        ) -> Tuple[nn.Module]:
+        """
+        Constructs a GRU 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.
+        Returns:
+            Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
+        """
+        with torch.inference_mode(False):
+            model = nn.GRU(
+                input_size=input_size,
+                hidden_size=hidden_size,
+                num_layers=num_layers,
+                bias=bias,
+                batch_first=batch_first,
+                dropout=dropout,
+                bidirectional=bidirectional
+            )
+
+            # TODO: Xaviar normal weight initialization is not
+            # set up for GRU. If performance turning is needed, consider adding it here.
+
+        return (model,)

modules/pytorch_wrapper/ptn_gru_linear.py

@@ -0,0 +1,163 @@
+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 GRULinear(nn.Module):
+    """
+    A recurrent neural network (GRU) with a linear output layer.
+    
+    This module uses a standard PyTorch GRU, 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,
+        linear_output_size: int,
+        linear_bias: bool
+    ):
+        super().__init__()
+        self.gru_model = nn.GRU(
+            input_size=input_size,
+            hidden_size=hidden_size,
+            num_layers=num_layers,
+            bias=bias,
+            batch_first=batch_first,
+            dropout=dropout,
+            bidirectional=bidirectional
+        )
+        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 GRU 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.gru_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 PtnGRULinear:
+    """
+    Ptn GRU Linear:
+    A recurrent neural network (GRU) 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 GRU.
+            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 GRU 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.
+            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}),
+                "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,
+          linear_output_size: int,
+          linear_bias: bool
+        ) -> Tuple[nn.Module]:
+        """
+        Constructs an GRU 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.
+            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 = GRULinear(
+                input_size=input_size,
+                hidden_size=hidden_size,
+                num_layers=num_layers,
+                bias=bias,
+                batch_first=batch_first,
+                dropout=dropout,
+                bidirectional=bidirectional,
+                linear_output_size=linear_output_size,
+                linear_bias=linear_bias
+            )
+            return (model,)

modules/pytorch_wrapper/ptn_rnn.py

@@ -105,12 +105,5 @@ class PtnRNN:
                         torch.nn.init.xavier_normal_(weight_ih)
                         torch.nn.init.xavier_normal_(weight_hh)
 
-                    # bias_ih = getattr(model, f"bias_ih_l{layer}{suffix}")
-                    # bias_hh = getattr(model, f"bias_hh_l{layer}{suffix}")
-
-                    # torch.nn.init.uniform_(weight_ih, -0.05, 0.05)
-                    # torch.nn.init.uniform_(weight_hh, -0.05, 0.05)
-                    # nn.init.zeros_(bias_ih)
-                    # nn.init.zeros_(bias_hh)
 
         return (model,)

pyproject.toml

@@ -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.0"
+version = "3.3.1"
 license = {file = "LICENSE"}
 
 [project.urls]

tests/pytorch_wrapper/pt_train_gru_linear_text_classification_model_test.py

@@ -0,0 +1,125 @@
+"""
+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_gru_linear import PtnGRULinear
+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 TestPtTrainGRULinearTextClassificationModel(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 = "gru_text_classification_test.pt"
+
+        self.model_node = PtnGRULinear()
+        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,
+            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()

tests/pytorch_wrapper/ptn_gru_linear_test.py

@@ -0,0 +1,87 @@
+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_gru_linear import PtnGRULinear
+
+
+class TestPtnGRULinear(unittest.TestCase):
+    def setUp(self):
+        """Set up test instance."""
+        self.model_node = PtnGRULinear()
+
+    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,
+            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,
+            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,
+            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,
+            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()

tests/pytorch_wrapper/ptn_gru_test.py

@@ -0,0 +1,148 @@
+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_gru import PtnGRU
+
+
+class TestPtnGRU(unittest.TestCase):
+    def setUp(self):
+        """Set up test instance."""
+        self.model_node = PtnGRU()
+
+    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]
+
+        x = torch.randn(10, 8, 2)  # [Seq, B, Token]
+        out_seq_last_layer, out_last_value = 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]
+
+        x = torch.randn(10, 8, 2)  # [Seq, B, Token]
+        out_seq_last_layer, out_last_value = 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]
+
+        x = torch.randn(10, 8, 2)  # [Seq, B, Token]
+        out_seq_last_layer, out_last_value = 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]
+
+        x = torch.randn(10, 8, 2)  # [Seq, B, Token]
+        out_seq_last_layer, out_last_value = 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]
+
+        x = torch.randn(8, 10, 2)  # [B, Seq, Token]
+        out_seq_last_layer, out_last_value = 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]
+
+        x = torch.randn(8, 10, 2)  # [B, Seq, Token]
+        out_seq_last_layer, out_last_value = 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]
+
+        x = torch.randn(8, 10, 2)  # [B, Seq, Token]
+        out_seq_last_layer, out_last_value = 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]
+
+        x = torch.randn(8, 10, 2)  # [B, Seq, Token]
+        out_seq_last_layer, out_last_value = 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()