diff --git a/README.md b/README.md index b2ba28c..1b5b2ce 100644 --- a/README.md +++ b/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. diff --git a/__init__.py b/__init__.py index ba29d23..9dc8235 100644 --- a/__init__.py +++ b/__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", diff --git a/dev/scripts/init_py_generate.py b/dev/scripts/init_py_generate.py index 72d94ce..241826f 100644 --- a/dev/scripts/init_py_generate.py +++ b/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"), diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index ab3bea4..d59de02 100644 --- a/docs/reference/node_reference.md +++ b/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. | diff --git a/docs/reference/pt_train_rnn_model.md b/docs/reference/pt_train_rnn_model.md index 9df7e61..901484e 100644 --- a/docs/reference/pt_train_rnn_model.md +++ b/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 | diff --git a/docs/reference/ptn_gru.md b/docs/reference/ptn_gru.md new file mode 100644 index 0000000..ee1cff6 --- /dev/null +++ b/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 | + +
+Category: PyTorch wrapper - Model + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptn_gru_linear.md b/docs/reference/ptn_gru_linear.md new file mode 100644 index 0000000..135115f --- /dev/null +++ b/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 | + +
+Category: PyTorch wrapper - Model + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptn_rnn_linear.md b/docs/reference/ptn_rnn_linear.md index 764670f..817925c 100644 --- a/docs/reference/ptn_rnn_linear.md +++ b/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 | diff --git a/docs/training_rnn_for_classification.md b/docs/training_rnn_for_classification.md index 808bdad..bccbe53 100644 --- a/docs/training_rnn_for_classification.md +++ b/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**. diff --git a/examples/workflows/gru_classification.json b/examples/workflows/gru_classification.json new file mode 100644 index 0000000..6531ffd --- /dev/null +++ b/examples/workflows/gru_classification.json @@ -0,0 +1 @@ +{"last_node_id":21,"last_link_id":26,"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":23},{"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/gru_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":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/gru_classification.pt"]},{"id":20,"type":"PtnGRULinear","pos":[2672.0244140625,1382.3995361328125],"size":[315,250],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[23,24],"slot_index":0}],"properties":{"Node name for S&R":"PtnGRULinear"},"widgets_values":[100,512,2,true,true,0.4,true,1,true]},{"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":24}],"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]}],"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"],[23,20,0,14,0,"PTMODEL"],[24,20,0,21,0,"PTMODEL"],[25,21,0,15,0,"PTOPTIMIZER"],[26,21,0,14,2,"PTOPTIMIZER"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.3972131526604054,"offset":[-634.2050513544009,-140.5171891315913]},"node_versions":{"pt-wrapper":"b4bd168592762ed6fdc7dc4261d1f5fdf0d91e30","data-analysis":"68dd090367adffa569c3b3993d56e422f4516d2d"},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/modules/pytorch_wrapper/hf_glove_dataset.py b/modules/pytorch_wrapper/hf_glove_dataset.py index 1fdf4ff..48c4768 100644 --- a/modules/pytorch_wrapper/hf_glove_dataset.py +++ b/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() diff --git a/modules/pytorch_wrapper/pt_train_rnn_model.py b/modules/pytorch_wrapper/pt_train_rnn_model.py index 153ca2d..558dff3 100644 --- a/modules/pytorch_wrapper/pt_train_rnn_model.py +++ b/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, diff --git a/modules/pytorch_wrapper/ptn_gru.py b/modules/pytorch_wrapper/ptn_gru.py new file mode 100644 index 0000000..437ce8f --- /dev/null +++ b/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,) diff --git a/modules/pytorch_wrapper/ptn_gru_linear.py b/modules/pytorch_wrapper/ptn_gru_linear.py new file mode 100644 index 0000000..69dd86a --- /dev/null +++ b/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,) diff --git a/modules/pytorch_wrapper/ptn_rnn.py b/modules/pytorch_wrapper/ptn_rnn.py index e1660bb..332e47e 100644 --- a/modules/pytorch_wrapper/ptn_rnn.py +++ b/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,) diff --git a/pyproject.toml b/pyproject.toml index 7d11441..65fc309 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.0" +version = "3.3.1" license = {file = "LICENSE"} [project.urls] diff --git a/tests/pytorch_wrapper/pt_train_gru_linear_text_classification_model_test.py b/tests/pytorch_wrapper/pt_train_gru_linear_text_classification_model_test.py new file mode 100644 index 0000000..dfc3402 --- /dev/null +++ b/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() diff --git a/tests/pytorch_wrapper/ptn_gru_linear_test.py b/tests/pytorch_wrapper/ptn_gru_linear_test.py new file mode 100644 index 0000000..2114772 --- /dev/null +++ b/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() diff --git a/tests/pytorch_wrapper/ptn_gru_test.py b/tests/pytorch_wrapper/ptn_gru_test.py new file mode 100644 index 0000000..f6a5675 --- /dev/null +++ b/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()