diff --git a/README.md b/README.md index 1b5b2ce..487c48b 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/__init__.py b/__init__.py index 9dc8235..03dcef8 100644 --- a/__init__.py +++ b/__init__.py @@ -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", diff --git a/dev/scripts/init_py_generate.py b/dev/scripts/init_py_generate.py index 241826f..db7d2d4 100644 --- a/dev/scripts/init_py_generate.py +++ b/dev/scripts/init_py_generate.py @@ -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"), diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index d59de02..cb276b0 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -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. | diff --git a/docs/reference/ptn_lstm.md b/docs/reference/ptn_lstm.md new file mode 100644 index 0000000..41b8326 --- /dev/null +++ b/docs/reference/ptn_lstm.md @@ -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 | + +
+Category: PyTorch wrapper - Model + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptn_lstm_linear.md b/docs/reference/ptn_lstm_linear.md new file mode 100644 index 0000000..dd274e5 --- /dev/null +++ b/docs/reference/ptn_lstm_linear.md @@ -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 | + +
+Category: PyTorch wrapper - Model + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/training_rnn_for_classification.md b/docs/training_rnn_for_classification.md index bccbe53..4def5b8 100644 --- a/docs/training_rnn_for_classification.md +++ b/docs/training_rnn_for_classification.md @@ -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**. diff --git a/examples/workflows/lstm_classification.json b/examples/workflows/lstm_classification.json new file mode 100644 index 0000000..0b1de88 --- /dev/null +++ b/examples/workflows/lstm_classification.json @@ -0,0 +1 @@ +{"last_node_id":22,"last_link_id":28,"nodes":[{"id":12,"type":"PtDataLoader","pos":[2290.658447265625,2016.2130126953125],"size":[271.68658447265625,173.3039093017578],"flags":{},"order":5,"mode":0,"inputs":[{"name":"dataset","localized_name":"dataset","type":"PTVDATASET","link":10}],"outputs":[{"name":"PTDATALOADER","localized_name":"PTDATALOADER","type":"PTDATALOADER","links":[14],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[128,false,"{\"num_workers\":1}"]},{"id":7,"type":"PtDataLoader","pos":[2293.7060546875,1735.56494140625],"size":[271.68658447265625,173.3039093017578],"flags":{},"order":4,"mode":0,"inputs":[{"name":"dataset","localized_name":"dataset","type":"PTVDATASET","link":11}],"outputs":[{"name":"PTDATALOADER","localized_name":"PTDATALOADER","type":"PTDATALOADER","links":[13],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[128,true,"{\"num_workers\":1}"]},{"id":14,"type":"PtTrainRNNModel","pos":[3111.27294921875,1735.56494140625],"size":[315,346],"flags":{},"order":8,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":28},{"name":"train_loader","localized_name":"train_loader","type":"PTDATALOADER","link":13},{"name":"optimizer","localized_name":"optimizer","type":"PTOPTIMIZER","link":26},{"name":"loss_function","localized_name":"loss_function","type":"PTLOSS","link":18},{"name":"scheduler","localized_name":"scheduler","type":"PTLRSCHEDULER","shape":7,"link":16},{"name":"val_loader","localized_name":"val_loader","type":"PTDATALOADER","shape":7,"link":14},{"name":"h_0","localized_name":"h_0","type":"TENSOR","shape":7,"link":null}],"outputs":[{"name":"Model","localized_name":"Model","type":"PTMODEL","links":[19],"slot_index":0},{"name":"train loss","localized_name":"train loss","type":"TENSOR","links":null},{"name":"val loss","localized_name":"val loss","type":"TENSOR","links":null}],"properties":{"Node name for S&R":"PtTrainRNNModel"},"widgets_values":[true,50,true,false,10,false,true,true]},{"id":16,"type":"PtnBCEWithLogitsLoss","pos":[2668.492919921875,1800.2408447265625],"size":[315,58],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[{"name":"PTLOSS","localized_name":"PTLOSS","type":"PTLOSS","links":[18],"slot_index":0}],"properties":{"Node name for S&R":"PtnBCEWithLogitsLoss"},"widgets_values":["mean"]},{"id":10,"type":"PtvHfGloveDataset","pos":[1927.5224609375,1735.56494140625],"size":[315,178],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[{"name":"PTVDATASET","localized_name":"PTVDATASET","type":"PTVDATASET","links":[11],"slot_index":0}],"properties":{"Node name for S&R":"PtvHfGloveDataset"},"widgets_values":["imdb","train",100,300,"text","label"]},{"id":11,"type":"PtvHfGloveDataset","pos":[1920.82177734375,2017.7349853515625],"size":[315,178],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[{"name":"PTVDATASET","localized_name":"PTVDATASET","type":"PTVDATASET","links":[10],"slot_index":0}],"properties":{"Node name for S&R":"PtvHfGloveDataset"},"widgets_values":["imdb","test",100,300,"text","label"]},{"id":19,"type":"PandasShowText","pos":[3504.62841796875,1897.03662109375],"size":[302.7994079589844,112],"flags":{},"order":10,"mode":0,"inputs":[{"name":"text","type":"STRING","widget":{"name":"text"},"link":20}],"outputs":[{"name":"STRING","localized_name":"STRING","type":"STRING","shape":6,"links":null}],"properties":{"Node name for S&R":"PandasShowText"},"widgets_values":["","/container/transfer/models/lstm_classification.pt"]},{"id":15,"type":"PtoLrSchedulerReduceOnPlateau","pos":[3111.27294921875,1377.6142578125],"size":[314.1559143066406,82],"flags":{},"order":7,"mode":0,"inputs":[{"name":"optimizer","localized_name":"optimizer","type":"PTOPTIMIZER","link":25}],"outputs":[{"name":"PTLRSCHEDULER","localized_name":"PTLRSCHEDULER","type":"PTLRSCHEDULER","links":[16],"slot_index":0}],"properties":{"Node name for S&R":"PtoLrSchedulerReduceOnPlateau"},"widgets_values":[10,0.5000000000000002]},{"id":21,"type":"PtoAdamW","pos":[2676.04931640625,1162.1456298828125],"size":[315,154],"flags":{},"order":6,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":27}],"outputs":[{"name":"PTOPTIMIZER","localized_name":"PTOPTIMIZER","type":"PTOPTIMIZER","links":[25,26],"slot_index":0}],"properties":{"Node name for S&R":"PtoAdamW"},"widgets_values":[0.001,0.9,0.999,0.01,false]},{"id":22,"type":"PtnLSTMLinear","pos":[2680.248779296875,1391.870849609375],"size":[315,274],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[27,28],"slot_index":0}],"properties":{"Node name for S&R":"PtnLSTMLinear"},"widgets_values":[100,512,2,true,true,0.4,true,0,1,true]},{"id":17,"type":"PtSaveModel","pos":[3495.453125,1735.56494140625],"size":[318.3275146484375,91.27499389648438],"flags":{},"order":9,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":19}],"outputs":[{"name":"model full path","localized_name":"model full path","type":"STRING","links":[20],"slot_index":0}],"properties":{"Node name for S&R":"PtSaveModel"},"widgets_values":["/container/transfer/models/lstm_classification.pt"]}],"links":[[10,11,0,12,0,"PTVDATASET"],[11,10,0,7,0,"PTVDATASET"],[13,7,0,14,1,"PTDATALOADER"],[14,12,0,14,5,"PTDATALOADER"],[16,15,0,14,4,"PTLRSCHEDULER"],[18,16,0,14,3,"PTLOSS"],[19,14,0,17,0,"PTMODEL"],[20,17,0,19,0,"STRING"],[25,21,0,15,0,"PTOPTIMIZER"],[26,21,0,14,2,"PTOPTIMIZER"],[27,22,0,21,0,"PTMODEL"],[28,22,0,14,0,"PTMODEL"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.6397157544911096,"offset":[-1591.995279947916,-928.1316784347933]},"node_versions":{"pt-wrapper":"70586c134bf52d39b4a3b7275d8312964699e2be","data-analysis":"68dd090367adffa569c3b3993d56e422f4516d2d"},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/modules/pytorch_wrapper/ptn_lstm.py b/modules/pytorch_wrapper/ptn_lstm.py new file mode 100644 index 0000000..589b1ac --- /dev/null +++ b/modules/pytorch_wrapper/ptn_lstm.py @@ -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,) diff --git a/modules/pytorch_wrapper/ptn_lstm_linear.py b/modules/pytorch_wrapper/ptn_lstm_linear.py new file mode 100644 index 0000000..fa109bd --- /dev/null +++ b/modules/pytorch_wrapper/ptn_lstm_linear.py @@ -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,) diff --git a/pyproject.toml b/pyproject.toml index 65fc309..4cb6c8b 100644 --- a/pyproject.toml +++ b/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.1" +version = "3.3.2" license = {file = "LICENSE"} [project.urls] diff --git a/tests/pytorch_wrapper/pt_train_lstm_linear_text_classification_model_test.py b/tests/pytorch_wrapper/pt_train_lstm_linear_text_classification_model_test.py new file mode 100644 index 0000000..ccefc89 --- /dev/null +++ b/tests/pytorch_wrapper/pt_train_lstm_linear_text_classification_model_test.py @@ -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() diff --git a/tests/pytorch_wrapper/ptn_lstm_linear_test.py b/tests/pytorch_wrapper/ptn_lstm_linear_test.py new file mode 100644 index 0000000..5a649aa --- /dev/null +++ b/tests/pytorch_wrapper/ptn_lstm_linear_test.py @@ -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() diff --git a/tests/pytorch_wrapper/ptn_lstm_test.py b/tests/pytorch_wrapper/ptn_lstm_test.py new file mode 100644 index 0000000..37624a7 --- /dev/null +++ b/tests/pytorch_wrapper/ptn_lstm_test.py @@ -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()