From 37d0f9da6297b7ca3ddcdb6281c14edd2c87e0eb Mon Sep 17 00:00:00 2001 From: HowToSD Date: Mon, 24 Mar 2025 23:21:35 -0700 Subject: [PATCH] Added support for RNN --- __init__.py | 15 + dev/scripts/init_py_generate.py | 1 + docs/reference/node_reference.md | 5 + docs/reference/pt_train_rnn_model.md | 48 ++++ docs/reference/ptn_rnn.md | 34 +++ docs/reference/ptn_rnn_linear.md | 36 +++ docs/reference/ptv_hf_glove_dataset.md | 30 ++ .../ptv_sequential_tensor_dataset.md | 22 ++ examples/workflows/rnn_classification.json | 1 + modules/pytorch_wrapper/hf_glove_dataset.py | 129 +++++++++ modules/pytorch_wrapper/pt_train_model.py | 3 +- modules/pytorch_wrapper/pt_train_rnn_model.py | 262 ++++++++++++++++++ modules/pytorch_wrapper/ptn_rnn.py | 116 ++++++++ modules/pytorch_wrapper/ptn_rnn_linear.py | 184 ++++++++++++ .../pytorch_wrapper/ptv_hf_glove_dataset.py | 80 ++++++ .../ptv_sequential_tensor_dataset.py | 99 +++++++ modules/pytorch_wrapper/rnn_output_reshape.py | 84 ++++++ modules/pytorch_wrapper/utils.py | 16 ++ pyproject.toml | 2 +- requirements.txt | 2 + .../pytorch_wrapper/hf_glove_dataset_test.py | 64 +++++ .../pt_train_rnn_linear_model_test.py | 97 +++++++ ...n_linear_text_classification_model_test.py | 122 ++++++++ .../pt_train_rnn_model_test.py | 96 +++++++ tests/pytorch_wrapper/ptn_rnn_linear_test.py | 91 ++++++ tests/pytorch_wrapper/ptn_rnn_test.py | 156 +++++++++++ .../ptv_hf_glove_dataset_test.py | 50 ++++ .../ptv_sequential_tensor_dataset_test.py | 83 ++++++ .../rnn_output_reshape_test.py | 152 ++++++++++ 29 files changed, 2078 insertions(+), 2 deletions(-) create mode 100644 docs/reference/pt_train_rnn_model.md create mode 100644 docs/reference/ptn_rnn.md create mode 100644 docs/reference/ptn_rnn_linear.md create mode 100644 docs/reference/ptv_hf_glove_dataset.md create mode 100644 docs/reference/ptv_sequential_tensor_dataset.md create mode 100644 examples/workflows/rnn_classification.json create mode 100644 modules/pytorch_wrapper/hf_glove_dataset.py create mode 100644 modules/pytorch_wrapper/pt_train_rnn_model.py create mode 100644 modules/pytorch_wrapper/ptn_rnn.py create mode 100644 modules/pytorch_wrapper/ptn_rnn_linear.py create mode 100644 modules/pytorch_wrapper/ptv_hf_glove_dataset.py create mode 100644 modules/pytorch_wrapper/ptv_sequential_tensor_dataset.py create mode 100644 modules/pytorch_wrapper/rnn_output_reshape.py create mode 100644 tests/pytorch_wrapper/hf_glove_dataset_test.py create mode 100644 tests/pytorch_wrapper/pt_train_rnn_linear_model_test.py create mode 100644 tests/pytorch_wrapper/pt_train_rnn_linear_text_classification_model_test.py create mode 100644 tests/pytorch_wrapper/pt_train_rnn_model_test.py create mode 100644 tests/pytorch_wrapper/ptn_rnn_linear_test.py create mode 100644 tests/pytorch_wrapper/ptn_rnn_test.py create mode 100644 tests/pytorch_wrapper/ptv_hf_glove_dataset_test.py create mode 100644 tests/pytorch_wrapper/ptv_sequential_tensor_dataset_test.py create mode 100644 tests/pytorch_wrapper/rnn_output_reshape_test.py diff --git a/__init__.py b/__init__.py index e40fc7b..ba29d23 100644 --- a/__init__.py +++ b/__init__.py @@ -111,6 +111,7 @@ from .modules.pytorch_wrapper.pt_train_classification_model import PtTrainClassi from .modules.pytorch_wrapper.pt_train_classification_model_lr import PtTrainClassificationModelLr from .modules.pytorch_wrapper.pt_train_model import PtTrainModel from .modules.pytorch_wrapper.pt_train_regression_model import PtTrainRegressionModel +from .modules.pytorch_wrapper.pt_train_rnn_model import PtTrainRNNModel from .modules.pytorch_wrapper.pt_unsqueeze import PtUnsqueeze from .modules.pytorch_wrapper.pt_var import PtVar from .modules.pytorch_wrapper.pt_view import PtView @@ -168,6 +169,8 @@ from .modules.pytorch_wrapper.ptn_nll_loss import PtnNLLLoss from .modules.pytorch_wrapper.ptn_pre_add_channel_axis import PtnPreAddChannelAxis from .modules.pytorch_wrapper.ptn_pre_flatten import PtnPreFlatten from .modules.pytorch_wrapper.ptn_resnet_model import PtnResnetModel +from .modules.pytorch_wrapper.ptn_rnn import PtnRNN +from .modules.pytorch_wrapper.ptn_rnn_linear import PtnRNNLinear from .modules.pytorch_wrapper.ptn_smooth_l1_loss import PtnSmoothL1Loss from .modules.pytorch_wrapper.pto_adam import PtoAdam from .modules.pytorch_wrapper.pto_adamw import PtoAdamW @@ -179,7 +182,9 @@ from .modules.pytorch_wrapper.pto_simple import PtoSimple from .modules.pytorch_wrapper.ptv_dataset import PtvDataset from .modules.pytorch_wrapper.ptv_dataset_len import PtvDatasetLen from .modules.pytorch_wrapper.ptv_dataset_loader import PtvDatasetLoader +from .modules.pytorch_wrapper.ptv_hf_glove_dataset import PtvHfGloveDataset from .modules.pytorch_wrapper.ptv_image_folder_dataset import PtvImageFolderDataset +from .modules.pytorch_wrapper.ptv_sequential_tensor_dataset import PtvSequentialTensorDataset from .modules.pytorch_wrapper.ptv_transforms_data_augment import PtvTransformsDataAugment from .modules.pytorch_wrapper.ptv_transforms_resize import PtvTransformsResize from .modules.pytorch_wrapper.ptv_transforms_to_tensor import PtvTransformsToTensor @@ -302,6 +307,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtTrainClassificationModel": PtTrainClassificationModel, "PtTrainClassificationModelLr": PtTrainClassificationModelLr, "PtTrainModel": PtTrainModel, + "PtTrainRNNModel": PtTrainRNNModel, "PtTrainRegressionModel": PtTrainRegressionModel, "PtUnsqueeze": PtUnsqueeze, "PtVar": PtVar, @@ -359,6 +365,8 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtnNLLLoss": PtnNLLLoss, "PtnPreAddChannelAxis": PtnPreAddChannelAxis, "PtnPreFlatten": PtnPreFlatten, + "PtnRNN": PtnRNN, + "PtnRNNLinear": PtnRNNLinear, "PtnResnetModel": PtnResnetModel, "PtnSmoothL1Loss": PtnSmoothL1Loss, "PtoAdam": PtoAdam, @@ -371,7 +379,9 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtvDataset": PtvDataset, "PtvDatasetLen": PtvDatasetLen, "PtvDatasetLoader": PtvDatasetLoader, + "PtvHfGloveDataset": PtvHfGloveDataset, "PtvImageFolderDataset": PtvImageFolderDataset, + "PtvSequentialTensorDataset": PtvSequentialTensorDataset, "PtvTransformsDataAugment": PtvTransformsDataAugment, "PtvTransformsResize": PtvTransformsResize, "PtvTransformsToTensor": PtvTransformsToTensor, @@ -494,6 +504,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtTrainClassificationModel": "Pt Train Classification Model", "PtTrainClassificationModelLr": "Pt Train Classification Model Lr", "PtTrainModel": "Pt Train Model", + "PtTrainRNNModel": "Pt Train RNN Model", "PtTrainRegressionModel": "Pt Train Regression Model", "PtUnsqueeze": "Pt Unsqueeze", "PtVar": "Pt Var", @@ -551,6 +562,8 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtnNLLLoss": "Ptn NLL Loss", "PtnPreAddChannelAxis": "Ptn Pre Add Channel Axis", "PtnPreFlatten": "Ptn Pre Flatten", + "PtnRNN": "Ptn RNN", + "PtnRNNLinear": "Ptn RNN Linear", "PtnResnetModel": "Ptn Resnet Model", "PtnSmoothL1Loss": "Ptn Smooth L1 Loss", "PtoAdam": "Pto Adam", @@ -563,7 +576,9 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtvDataset": "Ptv Dataset", "PtvDatasetLen": "Ptv Dataset Len", "PtvDatasetLoader": "Ptv Dataset Loader", + "PtvHfGloveDataset": "Ptv Hf Glove Dataset", "PtvImageFolderDataset": "Ptv Image Folder Dataset", + "PtvSequentialTensorDataset": "Ptv Sequential Tensor Dataset", "PtvTransformsDataAugment": "Ptv Transforms Data Augment", "PtvTransformsResize": "Ptv Transforms Resize", "PtvTransformsToTensor": "Ptv Transforms To Tensor", diff --git a/dev/scripts/init_py_generate.py b/dev/scripts/init_py_generate.py index 020794b..72d94ce 100644 --- a/dev/scripts/init_py_generate.py +++ b/dev/scripts/init_py_generate.py @@ -78,6 +78,7 @@ def format_display_name(class_name): ("MP L", "MPL"), ("SN S", "SNS"), (" N A", " NA"), + ("R N N", "RNN"), ("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 a050511..ab3bea4 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -162,6 +162,8 @@ | [Ptn Max Pool 2d](ptn_max_pool2d.md) | A maxpool layer. | | [Ptn Pre Add Channel Axis](ptn_pre_add_channel_axis.md) | Adds a channel axis after the batch axis if the input is rank 3 (bs, h, w) | | [Ptn Pre Flatten](ptn_pre_flatten.md) | Flattens the input tensor before processing the tensor in the specified model. | +| [Ptn RNN](ptn_rnn.md) | A recurrent neural network (RNN) model consisting of one or more of a recurrent layer. | +| [Ptn RNN Linear](ptn_rnn_linear.md) | A recurrent neural network (RNN) model with a linear head. | | [Ptn Resnet Model](ptn_resnet_model.md) | A Resnet model consisting of multiple Resnet layers. | ## PyTorch wrapper - Optimizer | Node | Description | @@ -226,6 +228,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 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. | @@ -233,7 +236,9 @@ | [Ptv Dataset](ptv_dataset.md) | A Torchvision Dataset class wrapper. | | [Ptv Dataset Len](ptv_dataset_len.md) | A wrapper class that calls Python len on a dataset. | | [Ptv Dataset Loader](ptv_dataset_loader.md) | A node to combine the dataset and data loader into a single node. | +| [Ptv Hf Glove Dataset](ptv_hf_glove_dataset.md) | A PyTorch Dataset class wrapper for a Hugging Face dataset that converts text to embedding using Glove. | | [Ptv Image Folder Dataset](ptv_image_folder_dataset.md) | A Torchvision ImageFolder Dataset class wrapper. | +| [Ptv Sequential Tensor Dataset](ptv_sequential_tensor_dataset.md) | Creates a sequential tensor Dataset. | | [Ptv Transforms Data Augment](ptv_transforms_data_augment.md) | Applies data augmentation transformations to dataset elements. | | [Ptv Transforms Resize](ptv_transforms_resize.md) | Resizes and transforms elements of dataset to PyTorch tensors. | | [Ptv Transforms To Tensor](ptv_transforms_to_tensor.md) | Transforms elements of dataset to PyTorch tensors. | diff --git a/docs/reference/pt_train_rnn_model.md b/docs/reference/pt_train_rnn_model.md new file mode 100644 index 0000000..9df7e61 --- /dev/null +++ b/docs/reference/pt_train_rnn_model.md @@ -0,0 +1,48 @@ +# Pt Train R N N Model +Trains an RNN model using a given dataset, loss function, optimizer, and number of epochs with learning rate decay. + + Args: + model (PTMODEL): The PyTorch model to be trained. + linear_head (bool): The RNN model has a Linear head. + train_loader (torch.utils.data.DataLoader): DataLoader containing the training dataset. + optimizer (PTOPTIMIZER): Optimizer used for training. + loss (PTLOSS): The loss class. + epochs (int): Number of training epochs. + use_gpu (bool): Whether to use GPU for training. + early_stopping (bool): Whether to turn on early stopping + early_stopping_rounds (int): Specifies the number of rounds to monitor the evaluation loss when early_stopping is enabled and validation data is specified. If the loss does not decrease compared to the current best loss within this period, training will be stopped. + output_best_val_model (bool): Whether to outputs the model with the best val loss or the model from the last epoch. + 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. + +## Input +| Name | Data type | +|---|---| +| model | Ptmodel | +| linear_head | Boolean | +| train_loader | Ptdataloader | +| optimizer | Ptoptimizer | +| loss_function | Ptloss | +| epochs | Int | +| use_gpu | Boolean | +| early_stopping | Boolean | +| early_stopping_rounds | Int | +| output_best_val_model | Boolean | +| classification_metrics | Boolean | +| use_valid_token_mean | Boolean | +| scheduler | Ptlrscheduler | +| val_loader | Ptdataloader | + +## Output +| Data type | +|---| +| Ptmodel | +| Tensor | +| Tensor | + +
+Category: PyTorch wrapper - Training + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptn_rnn.md b/docs/reference/ptn_rnn.md new file mode 100644 index 0000000..def9544 --- /dev/null +++ b/docs/reference/ptn_rnn.md @@ -0,0 +1,34 @@ +# Ptn RNN +A recurrent neural network (RNN) 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. + nonlinearity (str): Activation function to apply after each hidden layer. Specify 'tanh' or 'relu'. + 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 RNN 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 | +| nonlinearity | | +| 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_rnn_linear.md b/docs/reference/ptn_rnn_linear.md new file mode 100644 index 0000000..764670f --- /dev/null +++ b/docs/reference/ptn_rnn_linear.md @@ -0,0 +1,36 @@ +# Ptn RNN Linear +A recurrent neural network (RNN) 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 RNN. + nonlinearity (str): Activation function to apply after each hidden layer. Specify 'tanh' or 'relu'. + 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 RNN 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 | +| nonlinearity | | +| 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/ptv_hf_glove_dataset.md b/docs/reference/ptv_hf_glove_dataset.md new file mode 100644 index 0000000..812ffdd --- /dev/null +++ b/docs/reference/ptv_hf_glove_dataset.md @@ -0,0 +1,30 @@ +# Ptv Hf Glove Dataset +A PyTorch Dataset class wrapper for a Hugging Face dataset that converts text to embedding using Glove. + + Args: + name (str): Name of the dataset class (e.g., MNIST, FashionMNIST). + split (str): The name of the subset of dataset to retrieve or the split specifier. + embed_dim (int): Embedding dimension for Glove. + max_seq_len (int): Sequence length. + sample_field_name (str): Field name for text samples. + label_field_name (str): Field name for labels. + +## Input +| Name | Data type | +|---|---| +| name | String | +| split | String | +| embed_dim | Int | +| max_seq_len | Int | +| sample_field_name | String | +| label_field_name | String | + +## Output +| Data type | +|---| +| Ptvdataset | + +
+Category: PyTorch wrapper - Training + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptv_sequential_tensor_dataset.md b/docs/reference/ptv_sequential_tensor_dataset.md new file mode 100644 index 0000000..1dbcc86 --- /dev/null +++ b/docs/reference/ptv_sequential_tensor_dataset.md @@ -0,0 +1,22 @@ +# Ptv Sequential Tensor Dataset +Creates a sequential tensor Dataset. + + Args: + tens (torch.Tensor): The tensor containing the data. + seq_len (int): Sequence length. + +## Input +| Name | Data type | +|---|---| +| tens | Tensor | +| seq_len | Int | + +## Output +| Data type | +|---| +| Ptvdataset | + +
+Category: PyTorch wrapper - Training + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/examples/workflows/rnn_classification.json b/examples/workflows/rnn_classification.json new file mode 100644 index 0000000..7133602 --- /dev/null +++ b/examples/workflows/rnn_classification.json @@ -0,0 +1 @@ +{"last_node_id":19,"last_link_id":21,"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,326],"flags":{},"order":8,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":12},{"name":"train_loader","localized_name":"train_loader","type":"PTDATALOADER","link":13},{"name":"optimizer","localized_name":"optimizer","type":"PTOPTIMIZER","link":17},{"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}],"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":9,"type":"PtoAdam","pos":[2668.492919921875,1191.617431640625],"size":[315,106],"flags":{},"order":6,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":21}],"outputs":[{"name":"PTOPTIMIZER","localized_name":"PTOPTIMIZER","type":"PTOPTIMIZER","links":[15,17],"slot_index":0}],"properties":{"Node name for S&R":"PtoAdam"},"widgets_values":[0.001,0.9,0.999]},{"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":15}],"outputs":[{"name":"PTLRSCHEDULER","localized_name":"PTLRSCHEDULER","type":"PTLRSCHEDULER","links":[16],"slot_index":0}],"properties":{"Node name for S&R":"PtoLrSchedulerReduceOnPlateau"},"widgets_values":[3,0.5000000000000002]},{"id":19,"type":"PandasShowText","pos":[3504.62841796875,1897.03662109375],"size":[302.7994079589844,100],"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/rnn_classification.pt"]},{"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/rnn_classification.pt"]},{"id":13,"type":"PtnRNNLinear","pos":[2668.492919921875,1374.7393798828125],"size":[315,274],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[12,21],"slot_index":0}],"properties":{"Node name for S&R":"PtnRNNLinear"},"widgets_values":[100,128,2,"tanh",true,true,0.1,true,1,true]}],"links":[[10,11,0,12,0,"PTVDATASET"],[11,10,0,7,0,"PTVDATASET"],[12,13,0,14,0,"PTMODEL"],[13,7,0,14,1,"PTDATALOADER"],[14,12,0,14,5,"PTDATALOADER"],[15,9,0,15,0,"PTOPTIMIZER"],[16,15,0,14,4,"PTLRSCHEDULER"],[17,9,0,14,2,"PTOPTIMIZER"],[18,16,0,14,3,"PTLOSS"],[19,14,0,17,0,"PTMODEL"],[20,17,0,19,0,"STRING"],[21,13,0,9,0,"PTMODEL"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.7036873299402207,"offset":[-1643.4385821927863,-998.2899215112109]},"node_versions":{"pt-wrapper":"b4b5dc48f3843c4c802d035fac042d0d79bea5a1","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 new file mode 100644 index 0000000..1fdf4ff --- /dev/null +++ b/modules/pytorch_wrapper/hf_glove_dataset.py @@ -0,0 +1,129 @@ +from typing import Tuple +import re +from html import unescape +import torch +from torch.utils.data import Dataset +from datasets import load_dataset # Hugging Face (not PyTorch) +import scipy +import numpy as np +# Monkey-patch gensim as it relies on scipy's deprecated triu method +# /gensim/matutils.py has this line which causes ImportError: +# from scipy.linalg import get_blas_funcs, triu +# TODO: Monitor gensim and fix this if upgrade is made +scipy.linalg.triu = np.triu +import gensim.downloader as api + +class HfGloveDataset(Dataset): + """ + Hugging Face Dataset wrapper with word embedding support. + + pragma: skip_doc + """ + def __init__(self, + dataset_name: str, + split: str, + glove_dim: int = 100, + max_seq_len: int = 256, + sample_field_name: str = "text", + label_field_name: str = "label") -> None: + """ + Initializes the dataset by loading Hugging Face dataset and GloVe vectors. + + Args: + dataset_name (str): The name of the Hugging Face dataset. + split (str): The dataset split (e.g., 'train'). + glove_dim (int): Dimensionality of GloVe embeddings. + max_seq_len (int): Maximum sequence length. + sample_field_name (str): Field name for text samples. + label_field_name (str): Field name for labels. + """ + self.dataset = load_dataset(dataset_name, split=split) + self.glove_dim = glove_dim + self.max_seq_len = max_seq_len + self.sample_field_name = sample_field_name + self.label_field_name = label_field_name + + # Load GloVe from gensim + glove_name = f"glove-wiki-gigaword-{glove_dim}" + self.glove = api.load(glove_name) + self.unk_vector = torch.zeros(glove_dim) + + def tokenize(self, text: str) -> list[str]: + """ + Tokenizes raw text into a list of lowercase word tokens. + It removes HTML tags, decodes HTML entities, strips punctuation + (while keeping letters and numbers), and splits on whitespace. + + Args: + text (str): Raw text input (may include HTML and punctuation). + + Returns: + list[str]: List of clean, lowercase word tokens. + """ + # Decode HTML entities like "&", "<", etc. to their actual characters + text = unescape(text) + + # Remove HTML tags (like
,

, etc.) + # The regex "<.*?>" uses .*? (non-greedy match) to ensure it matches one tag at a time, + # instead of greedily removing everything between the first '<' and last '>' + text = re.sub(r"<.*?>", " ", text) + + # Remove all characters except letters, digits, and whitespace + # This strips punctuation like "!", ".", "(", but keeps words and numbers + text = re.sub(r"[^a-zA-Z0-9\s]", " ", text) + + # Convert to lowercase and split by whitespace to get token list + return text.lower().split() + + + def vectorize(self, tokens: list[str]) -> torch.Tensor: + """ + Converts tokens to a fixed-length tensor of word vectors. + + Args: + tokens (list[str]): A list of word tokens. + + Returns: + torch.Tensor: A tensor of shape (max_seq_len, glove_dim). + """ + vectors = [] + for token in tokens[:self.max_seq_len]: + if token in self.glove: + vec = torch.tensor(self.glove[token], dtype=torch.float32) + else: + vec = self.unk_vector + vectors.append(vec) + + if len(vectors) < self.max_seq_len: + pad = [torch.zeros(self.glove_dim) for _ in range(self.max_seq_len - len(vectors))] + vectors.extend(pad) + + return torch.stack(vectors) + + def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Retrieves the (embedded sequence, label) pair at the given index. + + Args: + idx (int): Index of the item. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: Embedding tensor and label tensor. + """ + # Get a text sample from database and convert to a word embedding + sample = self.dataset[idx] + tokens = self.tokenize(sample[self.sample_field_name]) + emb_tensor = self.vectorize(tokens) + + # Get a corresponding label + label = torch.tensor(sample[self.label_field_name], dtype=torch.long) + return emb_tensor, label + + def __len__(self) -> int: + """ + Returns the number of samples in the dataset. + + Returns: + int: Dataset length. + """ + return len(self.dataset) diff --git a/modules/pytorch_wrapper/pt_train_model.py b/modules/pytorch_wrapper/pt_train_model.py index c8735b0..67f7e18 100644 --- a/modules/pytorch_wrapper/pt_train_model.py +++ b/modules/pytorch_wrapper/pt_train_model.py @@ -65,7 +65,7 @@ class PtTrainModel: 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]: """ - Trains a classification model. + Trains a model. Args: model (torch.nn.Module): The PyTorch model to be trained. @@ -117,6 +117,7 @@ class PtTrainModel: labels = labels.to("cuda") optimizer.zero_grad() outputs = model(images) + loss = loss_function(outputs, labels) loss.backward() optimizer.step() diff --git a/modules/pytorch_wrapper/pt_train_rnn_model.py b/modules/pytorch_wrapper/pt_train_rnn_model.py new file mode 100644 index 0000000..153ca2d --- /dev/null +++ b/modules/pytorch_wrapper/pt_train_rnn_model.py @@ -0,0 +1,262 @@ +import os +import sys +from typing import Any, Dict, Optional, Tuple, Union +import torch +import torch.nn as nn +import torch.nn.functional as F +from sklearn.metrics import accuracy_score +from .rnn_output_reshape import extract_rnn_return_value_with_adjusted_label + +class PtTrainRNNModel: + """ + Trains an RNN model using a given dataset, loss function, optimizer, and number of epochs with learning rate decay. + + Args: + model (PTMODEL): The PyTorch model to be trained. + linear_head (bool): The RNN model has a Linear head. + train_loader (torch.utils.data.DataLoader): DataLoader containing the training dataset. + optimizer (PTOPTIMIZER): Optimizer used for training. + loss (PTLOSS): The loss class. + epochs (int): Number of training epochs. + use_gpu (bool): Whether to use GPU for training. + early_stopping (bool): Whether to turn on early stopping + early_stopping_rounds (int): Specifies the number of rounds to monitor the evaluation loss when early_stopping is enabled and validation data is specified. If the loss does not decrease compared to the current best loss within this period, training will be stopped. + output_best_val_model (bool): Whether to outputs the model with the best val loss or the model from the last epoch. + 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. + + category: PyTorch wrapper - Training + """ + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + """ + Defines the input types required for training the model. + + Returns: + Dict[str, Any]: A dictionary specifying required and optional input types. + """ + return { + "required": { + "model": ("PTMODEL", {}), + "linear_head": ("BOOLEAN", {"default":True}), + "train_loader": ("PTDATALOADER", {}), + "optimizer": ("PTOPTIMIZER", {}), + "loss_function": ("PTLOSS", {}), + "epochs": ("INT", {"default":1, "min":1, "max":1e6}), + "use_gpu": ("BOOLEAN", {"default":False}), + "early_stopping": ("BOOLEAN", {"default":False}), + "early_stopping_rounds": ("INT", {"default":10, "min":1, "max":1000}), + "output_best_val_model": ("BOOLEAN", {"default":True}), + "classification_metrics": ("BOOLEAN", {"default":False}), + "use_valid_token_mean" : ("BOOLEAN", {"default":True}) + }, + "optional": { + "scheduler": ("PTLRSCHEDULER", {}), + "val_loader": ("PTDATALOADER", {}), + } + } + + RETURN_NAMES: tuple = ("Model", "train loss", "val loss") + RETURN_TYPES: tuple = ("PTMODEL","TENSOR","TENSOR") + FUNCTION: str = "f" + CATEGORY: str = "Training" + + @classmethod + def IS_CHANGED(cls, **kw): + return float("NaN") + + def f(self, + model: nn.Module, + linear_head: bool, + train_loader: torch.utils.data.DataLoader, + optimizer: torch.optim.Optimizer, + loss_function: nn.Module, + epochs: int, + use_gpu: bool, + early_stopping: bool, + early_stopping_rounds: int, + output_best_val_model:bool, + 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]: + """ + Trains an RNN model. + + Args: + model (torch.nn.Module): The PyTorch model to be trained. + linear_head (bool): The RNN model has a Linear head. + train_loader (torch.utils.data.DataLoader): DataLoader containing the training dataset. + optimizer (torch.optim.Optimizer): Optimizer used for training. + loss (torch.nn.Module): The loss class. + epochs (int): Number of training epochs. + use_gpu (bool): Whether to use GPU for training. + early_stopping (bool): Whether to turn on early stopping + early_stopping_rounds (int): Specifies the number of rounds to monitor the evaluation loss when early_stopping is enabled and validation data is specified. If the loss does not decrease compared to the current best loss within this period, training will be stopped. + output_best_val_model (bool): Whether to outputs the model with the best val loss or the model from the last epoch. + 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 (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. + + Returns: + Tuple[torch.nn.Module, torch.Tensor, torch.Tensor]: A tuple containing the trained model, + training loss history, and validation loss history. + """ + if (scheduler is not None and + isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau) and + val_loader is None): + raise ValueError("Validation data loader needs to be specified when ReduceLROnPlateau scheduler is specified.") + + checkpoint_dir = os.path.realpath( + os.path.join(os.path.dirname(__file__), "..", "..", "models", "tmp")) + print(f"Checkpoint tmp directory is set to {checkpoint_dir}") + if os.path.exists(checkpoint_dir) is False: + os.makedirs(checkpoint_dir) + print(f"Created {checkpoint_dir}") + best_checkpoint_path = os.path.join(checkpoint_dir, "best_model.pt") + + best_val_loss = None + rounds_since_best_val_loss = 0 + + all_preds = [] + all_labels = [] + + with torch.inference_mode(False): + train_loss = [] + val_loss = [] + if use_gpu: + model.to("cuda") + + # Train the model + for epoch in range(epochs): + model.train() + total_loss = 0 + total_samples = len(train_loader.dataset) + for x, y in train_loader: + if use_gpu: + x = x.to("cuda") + y = y.to("cuda") + optimizer.zero_grad() + if linear_head: + 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) + outputs, adjusted_labels = extract_rnn_return_value_with_adjusted_label( + x, + y, + output_seq, + bidirectional=model.bidirectional, + return_valid_token_mean=use_valid_token_mean + ) + if outputs.size() != adjusted_labels.size(): + raise ValueError(f"y_hat size : {outputs.size()} does not match label size: {adjusted_labels.size()}") + y = adjusted_labels + y_hat = outputs + + loss = loss_function(y_hat, y) + loss.backward() + optimizer.step() + + batch_size = y.size(0) + batch_loss = loss.item() * batch_size + total_loss += batch_loss + + epoch_loss = total_loss / total_samples + print(f"Epoch (train) {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}") + train_loss.append(epoch_loss) + + if val_loader: + with torch.no_grad(): + model.train(False) # TODO: Change to model.*e*v*a*l() once Comfy's security checker is fixed. + total_loss = 0 + total_samples = len(val_loader.dataset) + for x, y in val_loader: + if use_gpu: + x = x.to("cuda") + y = y.to("cuda") + if linear_head: + 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) + outputs, adjusted_labels = extract_rnn_return_value_with_adjusted_label( + x, + y, + output_seq, + bidirectional=model.bidirectional, + return_valid_token_mean=use_valid_token_mean + ) + if outputs.size() != adjusted_labels.size(): + raise ValueError(f"y_hat size : {outputs.size()} does not match label size: {adjusted_labels.size()}") + y = adjusted_labels + y_hat = outputs + + loss = loss_function(y_hat, y) + batch_size = y.size(0) + batch_loss = loss.item() * batch_size + total_loss += batch_loss + + if classification_metrics and y.squeeze().dim() == 1 : + # Accuracy collection for binary classification + preds = (y_hat > 0.5).int().detach().cpu().numpy() + labels = y.int().detach().cpu().numpy() + all_preds.extend(preds) + all_labels.extend(labels) + + epoch_loss = total_loss / total_samples + print(f"Epoch (val) {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}") + val_loss.append(epoch_loss) + + if classification_metrics: + if len(all_labels): + # Compute accuracy + acc = accuracy_score(all_labels, all_preds) + print(f"Validation Accuracy: {acc:.4f}") + else: + print("Accuracy was not computed as the label is not scalar.") + + if best_val_loss is None or epoch_loss < best_val_loss: + best_val_loss = epoch_loss + rounds_since_best_val_loss = 0 + torch.save(model.state_dict(), best_checkpoint_path) + print(f"Updated best val loss. Saved the model as {best_checkpoint_path}") + else: + rounds_since_best_val_loss += 1 + if early_stopping: + if rounds_since_best_val_loss >= early_stopping_rounds: + print(f"Val loss did not improve in {early_stopping_rounds} rounds. Stopping training.") + break + + if scheduler: + if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau): + scheduler.step(epoch_loss) + else: + scheduler.step() + + if use_gpu: + model.to("cpu") + + if val_loader: + if epoch_loss > best_val_loss and output_best_val_model: + sd = torch.load(best_checkpoint_path, weights_only=True) + model.load_state_dict(sd) + print("Loaded the best model for output") + else: + print("Using the final epoch model for output.") + + return (model, + torch.tensor(train_loss, dtype=torch.float32), + torch.tensor(val_loss, dtype=torch.float32)) \ No newline at end of file diff --git a/modules/pytorch_wrapper/ptn_rnn.py b/modules/pytorch_wrapper/ptn_rnn.py new file mode 100644 index 0000000..e1660bb --- /dev/null +++ b/modules/pytorch_wrapper/ptn_rnn.py @@ -0,0 +1,116 @@ +from typing import Any, Dict, Tuple +import torch +import torch.nn as nn + + +class PtnRNN: + """ + Ptn RNN: + A recurrent neural network (RNN) 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. + nonlinearity (str): Activation function to apply after each hidden layer. Specify 'tanh' or 'relu'. + 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 RNN 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}), + "nonlinearity": (("tanh", "relu"),), + "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, + nonlinearity: str, + bias: bool, + batch_first: bool, + dropout: float, + bidirectional: bool + ) -> Tuple[nn.Module]: + """ + Constructs an RNN 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. + nonlinearity (str): Activation function to apply after each hidden layer. Specify 'tanh' or 'relu'. + 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.RNN( + input_size=input_size, + hidden_size=hidden_size, + num_layers=num_layers, + nonlinearity=nonlinearity, + bias=bias, + batch_first=batch_first, + dropout=dropout, + bidirectional=bidirectional + ) + + directions = 2 if bidirectional else 1 + + # Initialize weights + for layer in range(num_layers): + for direction in range(directions): + suffix = f"_reverse" if direction == 1 else "" + + weight_ih = getattr(model, f"weight_ih_l{layer}{suffix}") + weight_hh = getattr(model, f"weight_hh_l{layer}{suffix}") + + # ih for X, hh is for prev H to current H + if nonlinearity == "relu": + torch.nn.init.kaiming_normal_(weight_ih, nonlinearity='relu') + torch.nn.init.kaiming_normal_(weight_hh, nonlinearity='relu') + else: + 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/modules/pytorch_wrapper/ptn_rnn_linear.py b/modules/pytorch_wrapper/ptn_rnn_linear.py new file mode 100644 index 0000000..a1fb977 --- /dev/null +++ b/modules/pytorch_wrapper/ptn_rnn_linear.py @@ -0,0 +1,184 @@ +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 RNNLinear(nn.Module): + """ + A recurrent neural network (RNN) with a linear output layer. + + This module uses a standard PyTorch RNN, 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, + nonlinearity: str, + bias: bool, + batch_first: bool, + dropout: float, + bidirectional: bool, + linear_output_size: int, + linear_bias: bool + ): + super().__init__() + self.rnn_model = nn.RNN( + input_size=input_size, + hidden_size=hidden_size, + num_layers=num_layers, + nonlinearity=nonlinearity, + bias=bias, + batch_first=batch_first, + dropout=dropout, + bidirectional=bidirectional + ) + self.bidirectional = bidirectional + self.batch_first = batch_first + directions = 2 if bidirectional else 1 + + # Initialize RNN weights + for layer in range(num_layers): + for direction in range(directions): + suffix = f"_reverse" if direction == 1 else "" + weight_ih = getattr(self.rnn_model, f"weight_ih_l{layer}{suffix}") + weight_hh = getattr(self.rnn_model, f"weight_hh_l{layer}{suffix}") + + if nonlinearity == "relu": + torch.nn.init.kaiming_normal_(weight_ih, nonlinearity='relu') + torch.nn.init.kaiming_normal_(weight_hh, nonlinearity='relu') + else: + torch.nn.init.xavier_normal_(weight_ih) + torch.nn.init.xavier_normal_(weight_hh) + + 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 RNN 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.rnn_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 PtnRNNLinear: + """ + Ptn RNN Linear: + A recurrent neural network (RNN) 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 RNN. + nonlinearity (str): Activation function to apply after each hidden layer. Specify 'tanh' or 'relu'. + 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 RNN 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}), + "nonlinearity": (("tanh", "relu"),), + "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, + nonlinearity: str, + bias: bool, + batch_first: bool, + dropout: float, + bidirectional: bool, + linear_output_size: int, + linear_bias: bool + ) -> Tuple[nn.Module]: + """ + Constructs an RNN 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. + nonlinearity (str): Activation function to apply after each hidden layer. Specify 'tanh' or 'relu'. + 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 = RNNLinear( + input_size=input_size, + hidden_size=hidden_size, + num_layers=num_layers, + nonlinearity=nonlinearity, + 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/ptv_hf_glove_dataset.py b/modules/pytorch_wrapper/ptv_hf_glove_dataset.py new file mode 100644 index 0000000..a37f1e8 --- /dev/null +++ b/modules/pytorch_wrapper/ptv_hf_glove_dataset.py @@ -0,0 +1,80 @@ +from typing import Any, Dict, Tuple +import torch + + +class PtvHfGloveDataset: + """ + A PyTorch Dataset class wrapper for a Hugging Face dataset that converts text to embedding using Glove. + + Args: + name (str): Name of the dataset class (e.g., MNIST, FashionMNIST). + split (str): The name of the subset of dataset to retrieve or the split specifier. + embed_dim (int): Embedding dimension for Glove. + max_seq_len (int): Sequence length. + sample_field_name (str): Field name for text samples. + label_field_name (str): Field name for labels. + + category: PyTorch wrapper - Training + """ + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + """ + Defines the required input types. + + Returns: + Dict[str, Any]: A dictionary of required input types. + """ + return { + "required": { + "name": ("STRING", {"default": "", "multiline": False}), + "split": ("STRING", {"default": "train","multiline": False}), + "embed_dim": ("INT", {"default": 100, "min":32, "max": 4096}), + "max_seq_len": ("INT", {"default": 256, "min":8, "max": 4096}), + "sample_field_name": ("STRING", {"default": "text","multiline": False}), + "label_field_name": ("STRING", {"default": "label","multiline": False}), + } + } + + RETURN_TYPES: Tuple[str] = ("PTVDATASET",) + FUNCTION: str = "f" + CATEGORY: str = "Training" + + @classmethod + def IS_CHANGED(cls, **kw): + return float("NaN") + + def f(self, + name: str, + split: str, + embed_dim: int, + max_seq_len: int, + sample_field_name: str, + label_field_name: str) -> Tuple: + """ + Loads a dataset from Hugging Face with specified parameters. + + Args: + name (str): Name of the dataset class (e.g., MNIST, FashionMNIST). + split (str): The name of the subset of dataset to retrieve or the split specifier. + embed_dim (int): Embedding dimension for Glove. + max_seq_len (int): Sequence length. + sample_field_name (str): Field name for text samples. + label_field_name (str): Field name for labels. + + Returns: + Tuple: A tuple containing the dataset instance. + """ + # Keep import within this method so that relevant packages are imported + # only when they are actually needed. + from .hf_glove_dataset import HfGloveDataset + + with torch.inference_mode(False): + dc = HfGloveDataset( + dataset_name=name, + split=split, + glove_dim=embed_dim, + max_seq_len=max_seq_len, + sample_field_name=sample_field_name, + label_field_name=label_field_name) + return (dc,) diff --git a/modules/pytorch_wrapper/ptv_sequential_tensor_dataset.py b/modules/pytorch_wrapper/ptv_sequential_tensor_dataset.py new file mode 100644 index 0000000..83ef87d --- /dev/null +++ b/modules/pytorch_wrapper/ptv_sequential_tensor_dataset.py @@ -0,0 +1,99 @@ +from typing import Any, Dict, Optional, Callable, Tuple +import torch +from torch.utils.data import Dataset +import torch +from typing import Tuple + +class SequentialTensorDataset(Dataset): + """ + Implementation of a sequential dataset. + + pragma: skip_doc + """ + def __init__(self, tens: torch.Tensor, sequence_length: int): + """ + Args: + tens (Tensor): Input tensor of shape [N, ...] containing the full sequence. + sequence_length (int): Length of each input sequence window. + """ + if tens.dim() == 1: + self.data_tens = torch.unsqueeze(tens, -1) + else: + self.data_tens = tens + + self.sequence_length = sequence_length + + def __len__(self) -> int: + """ + Returns: + int: Number of valid (input, target) sequence pairs. + """ + return self.data_tens.size(0) - self.sequence_length - 1 + + def __getitem__(self, i: int) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Args: + i (int): Index of the sequence window. + + Returns: + Tuple[Tensor, Tensor]: A tuple containing: + - x: Input sequence of shape [sequence_length, ...] + - y: Target sequence of shape [sequence_length, ...], shifted by one step + """ + x = self.data_tens[i : i + self.sequence_length ] + y = self.data_tens[i + 1 : i + self.sequence_length + 1] + return x, y + + + +class PtvSequentialTensorDataset: + """ + Creates a sequential tensor Dataset. + + Args: + tens (torch.Tensor): The tensor containing the data. + seq_len (int): Sequence length. + + category: PyTorch wrapper - Training + """ + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + """ + Defines the required input types. + + Returns: + Dict[str, Any]: A dictionary of required input types. + """ + return { + "required": { + "tens": ("TENSOR", {}), + "seq_len": ("INT", {"default":32, "min":1, "max":1e5}) + } + } + + RETURN_TYPES: Tuple[str] = ("PTVDATASET",) + FUNCTION: str = "f" + CATEGORY: str = "Training" + + @classmethod + def IS_CHANGED(cls, **kw): + return float("NaN") + + def f(self, + tens: torch.Tensor, + seq_len: int) -> Tuple: + """ + Creates a sequential tensor Dataset. + + Args: + tens (torch.Tensor): The tensor containing the data. + seq_len (int): Sequence length. + + Returns: + Tuple: A tuple containing the dataset instance. + """ + + with torch.inference_mode(False): + dc = SequentialTensorDataset(tens, seq_len) + return (dc,) diff --git a/modules/pytorch_wrapper/rnn_output_reshape.py b/modules/pytorch_wrapper/rnn_output_reshape.py new file mode 100644 index 0000000..27c2fbe --- /dev/null +++ b/modules/pytorch_wrapper/rnn_output_reshape.py @@ -0,0 +1,84 @@ +# Return Value of RNN Model's forward() +""" +model = nn.RNN(...) # or nn.LSTM, nn.GRU +out_seq, _ = model(input) + +out_seq has shape: + (batch, seq_len, hidden_size) if batch_first=True + +To extract the output from the last time step of the top RNN layer: + +Unidirectional Case: + final_output = out_seq[:, -1, :] # shape: (batch, hidden_size) + +Bidirectional Case: + If the model is bidirectional, hidden_size in out_seq is already doubled. + No need to manually concatenate forward/backward states. + + final_output = out_seq[:, -1, :] # shape: (batch, 2 * hidden_size) + +To match label shape: + - If the label is 3D: shape (batch, seq_len, hidden_size), use label[:, -1, :] + - If bidirectional: repeat the last axis +""" +from typing import Tuple, Union +import torch + +def extract_rnn_return_value_with_adjusted_label( + inputs: torch.Tensor, + label: torch.Tensor, + rnn_output: torch.Tensor, + bidirectional: bool, + batch_first:bool = True, + process_label:bool=True, + return_valid_token_mean:bool=True +) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """ + Extracts the final output at the last time step from the top RNN layer (out_seq[:, -1, :]) + and adjusts the label shape to match it. + + Args: + inputs (torch.Tensor): The input (feature) tensor. + label (torch.Tensor): Ground truth tensor of shape (batch, seq_len, hidden_size). + rnn_output (torch.Tensor): The output sequence from the RNN (out_seq), + shape (batch, seq_len, hidden_size) for unidirectional, + or (batch, seq_len, 2 * hidden_size) for bidirectional RNNs. + bidirectional (bool): Whether the RNN is bidirectional. + batch_first (bool): True if the first axis is the batch. + process_label (bool): True if process label. To process only the RNN output, pass None to label and set this flag to False. + return_valid_token_mean (bool): True to return the mean of non-zero outputs. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: A tuple containing: + - The final time-step output of shape (batch, hidden_size) or (batch, 2 * hidden_size) + - The adjusted label of matching shape + """ + if batch_first: + if return_valid_token_mean: + # Use inputs to build a valid token mask + valid_mask = inputs.norm(p=2, dim=2) > 0 # shape: [batch, seq_len] + lengths = valid_mask.sum(dim=1).clamp(min=1) # shape: [batch] + masked = rnn_output * valid_mask.unsqueeze(-1) # [batch, seq_len, hidden] + out = masked.sum(dim=1) / lengths.unsqueeze(1) # [batch, hidden] + else: + out = rnn_output[:, -1, :] # last time step output from top layer + if process_label: + label = label[:, -1, :] + else: + if return_valid_token_mean: + valid_mask = inputs.norm(p=2, dim=2) > 0 # shape: [seq_len, batch] + lengths = valid_mask.sum(dim=0).clamp(min=1) # shape: [batch] + masked = rnn_output * valid_mask.unsqueeze(-1) # [seq_len, batch, hidden] + out = masked.sum(dim=0) / lengths.unsqueeze(1) # [batch, hidden] + else: + out = rnn_output[-1, :, :] # last time step output from top layer + if process_label: + label = label[-1, :, :] + + if bidirectional and process_label: + label = torch.cat([label, label], dim=-1) + + if process_label: + return out, label + else: + return out diff --git a/modules/pytorch_wrapper/utils.py b/modules/pytorch_wrapper/utils.py index d86a661..9c5bf5d 100644 --- a/modules/pytorch_wrapper/utils.py +++ b/modules/pytorch_wrapper/utils.py @@ -1,6 +1,8 @@ import os import ast from typing import Optional, Tuple, Union +import random +import numpy as np import torch @@ -137,3 +139,17 @@ def str_to_dim(dim: str) -> Optional[Union[int, Tuple[int, ...]]]: return None +def set_seed(seed=42) -> None: + """ + Sets seed for random number generation. + + Args: + seed (int): Initial seed. + """ + torch.manual_seed(seed) + np.random.seed(seed) + random.seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 365f883..7d11441 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.2.1" +version = "3.3.0" license = {file = "LICENSE"} [project.urls] diff --git a/requirements.txt b/requirements.txt index 42b57c1..f308048 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,3 +3,5 @@ pandas>=2.2.3 seaborn>=0.13.2 scipy>=1.14.1 scikit-learn>=1.6.1 +datasets>=3.4.1 +gensim==4.3.2 diff --git a/tests/pytorch_wrapper/hf_glove_dataset_test.py b/tests/pytorch_wrapper/hf_glove_dataset_test.py new file mode 100644 index 0000000..fdee69d --- /dev/null +++ b/tests/pytorch_wrapper/hf_glove_dataset_test.py @@ -0,0 +1,64 @@ +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.hf_glove_dataset import HfGloveDataset + + +class TestHfGloveDataset(unittest.TestCase): + """ + Tests Hugging Face Dataset wrapper with word embedding support. + """ + + def setUp(self): + """Set up test instance.""" + self.embedding_dim = 100 + self.max_seq_len = 256 + self.ds = HfGloveDataset( + "imdb", # dataset_name + "train", #split + self.embedding_dim, + self.max_seq_len + ) + + def test_vocabulary_size(self): + self.assertEqual(400000, len(self.ds.glove)) + + def test_tokenize(self): + """Test tokenization""" + s = "Maddie is a cute puppy." + actual = self.ds.tokenize(s) + expected = ['maddie', 'is', 'a', 'cute', 'puppy'] + self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.") + + def test_vectorize(self): + """Test embedding generation""" + s = "Maddie is a cute puppy." + tokens = self.ds.tokenize(s) + embeddings = self.ds.vectorize(tokens) + actual = embeddings.size() + expected = torch.Size((self.max_seq_len, self.embedding_dim)) + self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.") + + def test_loading(self): + it = iter(self.ds) + sample = next(it) + + # Compare text + actual = sample[0].size() + expected = torch.Size((self.max_seq_len, self.embedding_dim)) + self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.") + + # Compare label + actual = sample[1] + expected = torch.tensor(0, dtype=torch.int64) + self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/pt_train_rnn_linear_model_test.py b/tests/pytorch_wrapper/pt_train_rnn_linear_model_test.py new file mode 100644 index 0000000..1675339 --- /dev/null +++ b/tests/pytorch_wrapper/pt_train_rnn_linear_model_test.py @@ -0,0 +1,97 @@ +import os +import sys +import unittest +import torch +import torch.nn as nn + +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_rnn_linear import PtnRNNLinear +from pytorch_wrapper.pto_adam import PtoAdam +from pytorch_wrapper.ptv_sequential_tensor_dataset import PtvSequentialTensorDataset +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_step import PtoLrSchedulerStep +from pytorch_wrapper.ptn_mse_loss import PtnMSELoss +from pytorch_wrapper.utils import set_seed + + +class TestPtTrainRNNLinearModel(unittest.TestCase): + def setUp(self): + set_seed(49) + + self.loss_fun = PtnMSELoss().f(reduction="mean")[0] + dataset_node = PtvSequentialTensorDataset() + dataset = dataset_node.f( + torch.arange(0.01, 1.0, 0.01, dtype=torch.float32), + seq_len=4)[0] + data_loader_node = PtDataLoader() + + self.train_loader = data_loader_node.f( + dataset=dataset, + batch_size=128, + shuffle=True, + parameters='{"num_workers":1}' + )[0] + + self.model_node = PtnRNNLinear() + self.model = self.model_node.f( + 1, # input_size + 32, # hidden_size + 4, # num_layers + "tanh", # nonlinearity + True, # bias, + True, # batch_first [B, Seq, Token] + 0.0, # dropout + False, # bidirectional, + 1, # linear output size + True, # linear bias + )[0] + + self.optimizer = PtoAdam().f(self.model, 0.001, 0.9, 0.999)[0] + + self.scheduler = PtoLrSchedulerStep().f(self.optimizer, + 100, # steps + 0.1 # gamma + )[0] + + self.trainer = PtTrainRNNModel() + self.save_model = PtSaveModel() + self.load_model = PtLoadModel() + self.linear_head = True + self.epochs = 200 + + 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, + classification_metrics=False, + use_valid_token_mean=False, + output_best_val_model=False)[0] + self.save_model.f(trained_model, f"rnn_{self.epochs}_epochs.pt") + + def test_2(self): + """ + Tests prediction. Note that prediction can be off even if loss is close to 0. + """ + loaded_model = self.load_model.f(self.model, f"rnn_{self.epochs}_epochs.pt" )[0] + x = torch.tensor([[0.01, 0.02, 0.03, 0.04]], dtype=torch.float32).unsqueeze(-1) + y_hat = loaded_model(x) + # print(y_hat) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/pt_train_rnn_linear_text_classification_model_test.py b/tests/pytorch_wrapper/pt_train_rnn_linear_text_classification_model_test.py new file mode 100644 index 0000000..61e2c27 --- /dev/null +++ b/tests/pytorch_wrapper/pt_train_rnn_linear_text_classification_model_test.py @@ -0,0 +1,122 @@ +""" +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_rnn_linear import PtnRNNLinear +from pytorch_wrapper.pto_adam import PtoAdam +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 TestPtTrainRNNLinearTextClassificationModel(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 = "rnn_text_classification_test.pt" + + self.model_node = PtnRNNLinear() + self.model = self.model_node.f( + self.embedding_dim, # input_size + 128, # hidden_size + 2, # num_layers + "tanh", # nonlinearity + True, # bias, + True, # batch_first [B, Seq, Token] + 0.1, # dropout + True, # bidirectional, + 1, # linear output size + True, # linear bias + )[0] + + self.optimizer = PtoAdam().f(self.model, 0.001, 0.9, 0.999)[0] + + self.scheduler = PtoLrSchedulerReduceOnPlateau().f(self.optimizer, + grace_period=3, + 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/pt_train_rnn_model_test.py b/tests/pytorch_wrapper/pt_train_rnn_model_test.py new file mode 100644 index 0000000..a5c7581 --- /dev/null +++ b/tests/pytorch_wrapper/pt_train_rnn_model_test.py @@ -0,0 +1,96 @@ +import os +import sys +import unittest +import torch +import torch.nn as nn + +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_rnn import PtnRNN +from pytorch_wrapper.pto_adam import PtoAdam +from pytorch_wrapper.ptv_sequential_tensor_dataset import PtvSequentialTensorDataset +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_step import PtoLrSchedulerStep +from pytorch_wrapper.ptn_mse_loss import PtnMSELoss +from pytorch_wrapper.utils import set_seed + + +class TestPtTrainRNNModel(unittest.TestCase): + def setUp(self): + set_seed(43) + + self.loss_fun = PtnMSELoss().f(reduction="mean")[0] + dataset_node = PtvSequentialTensorDataset() + dataset = dataset_node.f( + torch.arange(1, 10, dtype=torch.float32) / 100.0, + 4)[0] + data_loader_node = PtDataLoader() + + self.train_loader = data_loader_node.f( + dataset=dataset, + batch_size=32, + shuffle=True, + parameters='{"num_workers":1}' + )[0] + + self.model_node = PtnRNN() + self.model = self.model_node.f( + 1, # input_size + 1, # hidden_size + 1, # num_layers + "tanh", # nonlinearity + True, # bias, + True, # batch_first [B, Seq, Token] + 0.0, # dropout + False # bidirectional + )[0] + + self.optimizer = PtoAdam().f(self.model, 0.1, 0.9, 0.999)[0] + + self.scheduler = PtoLrSchedulerStep().f(self.optimizer, + 100, # steps + 0.1 # gamma + )[0] + + self.trainer = PtTrainRNNModel() + self.save_model = PtSaveModel() + self.load_model = PtLoadModel() + self.epochs = 50 + self.linear_head = False + + 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=False, + use_valid_token_mean=True)[0] + self.save_model.f(trained_model, f"rnn_{self.epochs}_epochs.pt") + + def test_2(self): + """ + Tests prediction. Note that due to a hidden size set to 1, the model sometimes fails to + learn the sequence depending on the initial random weight values, so you may + see the output that is far off from 0.05 here. + """ + loaded_model = self.load_model.f(self.model, f"rnn_{self.epochs}_epochs.pt" )[0] + x = (torch.tensor([[1, 2, 3, 4]], dtype=torch.float32) / 100).unsqueeze(-1) + y_hat = loaded_model(x) + # print(y_hat) + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptn_rnn_linear_test.py b/tests/pytorch_wrapper/ptn_rnn_linear_test.py new file mode 100644 index 0000000..8c1c7d1 --- /dev/null +++ b/tests/pytorch_wrapper/ptn_rnn_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_rnn_linear import PtnRNNLinear + + +class TestPtnRNNLinear(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.model_node = PtnRNNLinear() + + def test_seq_first_1_layer_uni(self): + self.model = self.model_node.f( + 2, # input_size + 3, # hidden_size + 1, # num_layers + "tanh", # nonlinearity + 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 + "tanh", # nonlinearity + 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 + "tanh", # nonlinearity + 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 + "tanh", # nonlinearity + 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_rnn_test.py b/tests/pytorch_wrapper/ptn_rnn_test.py new file mode 100644 index 0000000..9dedc05 --- /dev/null +++ b/tests/pytorch_wrapper/ptn_rnn_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_rnn import PtnRNN + + +class TestPtnRNN(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.model_node = PtnRNN() + + def test_shape_1_layer(self): + self.model = self.model_node.f( + 2, # input_size + 3, # hidden_size + 1, # num_layers + "tanh", # nonlinearity + 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 + "tanh", # nonlinearity + 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 + "tanh", # nonlinearity + 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 + "tanh", # nonlinearity + 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 + "tanh", # nonlinearity + 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 + "tanh", # nonlinearity + 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 + "tanh", # nonlinearity + 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 + "tanh", # nonlinearity + 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() diff --git a/tests/pytorch_wrapper/ptv_hf_glove_dataset_test.py b/tests/pytorch_wrapper/ptv_hf_glove_dataset_test.py new file mode 100644 index 0000000..8a1959a --- /dev/null +++ b/tests/pytorch_wrapper/ptv_hf_glove_dataset_test.py @@ -0,0 +1,50 @@ +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.ptv_hf_glove_dataset import PtvHfGloveDataset + + +class TestPtvHfGloveDataset(unittest.TestCase): + """ + Tests the node for Hugging Face Dataset wrapper with word embedding support. + """ + + def setUp(self): + """Set up test instance.""" + self.embedding_dim = 100 + self.max_seq_len = 256 + self.node = PtvHfGloveDataset() + self.ds = self.node.f( + "imdb", # dataset_name + "train", #split + self.embedding_dim, + self.max_seq_len, + "text", + "label" + )[0] + + def test_tokenize(self): + """Test tokenization""" + s = "Maddie is a cute puppy." + actual = self.ds.tokenize(s) + expected = ['maddie', 'is', 'a', 'cute', 'puppy'] + self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.") + + def test_vectorize(self): + """Test embedding generation""" + s = "Maddie is a cute puppy." + tokens = self.ds.tokenize(s) + embeddings = self.ds.vectorize(tokens) + actual = embeddings.size() + expected = torch.Size((self.max_seq_len, self.embedding_dim)) + self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptv_sequential_tensor_dataset_test.py b/tests/pytorch_wrapper/ptv_sequential_tensor_dataset_test.py new file mode 100644 index 0000000..d976e38 --- /dev/null +++ b/tests/pytorch_wrapper/ptv_sequential_tensor_dataset_test.py @@ -0,0 +1,83 @@ +import os +import sys +import torch +import unittest + +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.ptv_sequential_tensor_dataset import PtvSequentialTensorDataset +from pytorch_wrapper.pt_data_loader import PtDataLoader + + +class TestPtvSequentialTensorDataset(unittest.TestCase): + def setUp(self): + """Set up test data.""" + self.node = PtvSequentialTensorDataset() + self.data = torch.arange(1, 100, dtype=torch.float32) + + def test_1(self): + """Tests first seq """ + d = self.node.f(self.data, seq_len=5)[0] + actual = next(iter(d))[0] + expected = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32).unsqueeze(-1) + self.assertTrue( + torch.allclose(actual, expected, atol=4), + f"Mismatch: actual: {actual}, expected: {expected}") + + def test_2(self): + """Tests last seq """ + d = self.node.f(self.data, seq_len=5)[0] + it = iter(d) + for _ in range(94): + next(it) + + actual = next(it)[0] + + # Get the tensor for the last seq + expected = torch.tensor([95, 96, 97, 98, 99], dtype=torch.float32).unsqueeze(-1) + self.assertEqual(actual.size(), expected.size()) + self.assertTrue( + torch.allclose(actual, expected, atol=4), + f"Mismatch: actual: {actual}, expected: {expected}") + + def test_3(self): + """Tests last seq """ + self.data = torch.arange(0.01, 1, 0.01, dtype=torch.float32) + d = self.node.f(self.data, seq_len=5)[0] + it = iter(d) + for _ in range(94): + next(it) + + actual = next(it)[0] + + # Get the tensor for the last seq + expected = torch.tensor([0.95, 0.96, 0.97, 0.98, 0.99], dtype=torch.float32).unsqueeze(-1) + self.assertEqual(actual.size(), expected.size()) + self.assertTrue( + torch.allclose(actual, expected, atol=4), + f"Mismatch: actual: {actual}, expected: {expected}") + + def test_4_loader(self): + """Tests loader """ + bs = 32 + seq_len = 5 + self.data = torch.arange(0.01, 1, 0.01, dtype=torch.float32) + dataset = self.node.f(self.data, seq_len=seq_len)[0] + data_loader_node = PtDataLoader() + + data_loader = data_loader_node.f( + dataset=dataset, + batch_size=bs, + shuffle=True, + parameters='{"num_workers":1}' + )[0] + + it = iter(data_loader) + x_batch, y_batch = next(it) + self.assertTrue(x_batch.size() == (bs, seq_len, 1)) + self.assertTrue(y_batch.size() == (bs, seq_len, 1)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/rnn_output_reshape_test.py b/tests/pytorch_wrapper/rnn_output_reshape_test.py new file mode 100644 index 0000000..192afd8 --- /dev/null +++ b/tests/pytorch_wrapper/rnn_output_reshape_test.py @@ -0,0 +1,152 @@ +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.rnn_output_reshape import extract_rnn_return_value_with_adjusted_label +from torch.nn import RNN + +class TestExtractRnnReturnValueWithAdjustedLabel(unittest.TestCase): + def run_test(self, num_layers, bidirectional, batch_first): + input_size = 2 + hidden_size = 3 + batch_size = 8 + seq_len = 10 + + model = RNN( + input_size=input_size, + hidden_size=hidden_size, + num_layers=num_layers, + nonlinearity="tanh", + bias=True, + batch_first=batch_first, + dropout=0.0, + bidirectional=bidirectional + ) + + if batch_first: + x = torch.randn(batch_size, seq_len, input_size) # [B, Seq, Token] + else: + x = torch.randn(seq_len, batch_size, input_size) # [Seq, B, Token] + + output_seq, _ = model(x) + + if batch_first: + label = torch.randn(batch_size, seq_len, hidden_size) + else: + label = torch.randn(seq_len, batch_size, hidden_size) + + out, adjusted_label = extract_rnn_return_value_with_adjusted_label( + x, + label, output_seq, bidirectional, batch_first=batch_first + ) + + expected_hidden_dim = hidden_size * (2 if bidirectional else 1) + self.assertEqual(out.shape, (batch_size, expected_hidden_dim)) + + self.assertEqual(adjusted_label.shape, (batch_size, expected_hidden_dim)) + + # y_hat and y + self.assertEqual(out.shape, adjusted_label.shape) + + def test_uni_1_layer_seq_first(self): + self.run_test(num_layers=1, bidirectional=False, batch_first=False) + + def test_uni_2_layer_seq_first(self): + self.run_test(num_layers=2, bidirectional=False, batch_first=False) + + def test_bi_1_layer_seq_first(self): + self.run_test(num_layers=1, bidirectional=True, batch_first=False) + + def test_bi_2_layer_seq_first(self): + self.run_test(num_layers=2, bidirectional=True, batch_first=False) + + def test_uni_1_layer_batch_first(self): + self.run_test(num_layers=1, bidirectional=False, batch_first=True) + + def test_uni_2_layer_batch_first(self): + self.run_test(num_layers=2, bidirectional=False, batch_first=True) + + def test_bi_1_layer_batch_first(self): + self.run_test(num_layers=1, bidirectional=True, batch_first=True) + + def test_bi_2_layer_batch_first(self): + self.run_test(num_layers=2, bidirectional=True, batch_first=True) + + def test_valid_token_mean_batch_first(self): + batch_size = 4 + seq_len = 6 + input_size = 5 + hidden_size = 7 + + rnn = RNN( + input_size=input_size, + hidden_size=hidden_size, + batch_first=True, + bidirectional=False + ) + + # input: two valid tokens followed by padding (zero vectors) + x = torch.cat([ + torch.randn(batch_size, 2, input_size), + torch.zeros(batch_size, seq_len - 2, input_size) + ], dim=1) + + label = torch.randn(batch_size, seq_len, hidden_size) + output_seq, _ = rnn(x) + + out, adjusted_label = extract_rnn_return_value_with_adjusted_label( + x, + label, + output_seq, + bidirectional=False, + batch_first=True, + process_label=True, + return_valid_token_mean=True + ) + + self.assertEqual(out.shape, (batch_size, hidden_size)) + self.assertEqual(adjusted_label.shape, (batch_size, hidden_size)) + + def test_valid_token_mean_seq_first(self): + batch_size = 4 + seq_len = 6 + input_size = 5 + hidden_size = 7 + + rnn = RNN( + input_size=input_size, + hidden_size=hidden_size, + batch_first=False, + bidirectional=True + ) + + # input: three valid tokens followed by padding (zero vectors) + x = torch.cat([ + torch.randn(3, batch_size, input_size), + torch.zeros(seq_len - 3, batch_size, input_size) + ], dim=0) + + label = torch.randn(seq_len, batch_size, hidden_size) + output_seq, _ = rnn(x) + + out, adjusted_label = extract_rnn_return_value_with_adjusted_label( + x, + label, + output_seq, + bidirectional=True, + batch_first=False, + process_label=True, + return_valid_token_mean=True + ) + + self.assertEqual(out.shape, (batch_size, hidden_size * 2)) + self.assertEqual(adjusted_label.shape, (batch_size, hidden_size * 2)) + + +if __name__ == "__main__": + unittest.main()