nodes/HowToSD-ComfyUI-Pt-Wrapper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added support for RNN

Browse Source
This commit is contained in:
HowToSD
2025-03-24 23:21:35 -07:00
parent b700c27160
commit 37d0f9da62
29 changed files with 2078 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+15
View File
@@ -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",
dev/scripts/init_py_generate.py
+1
View File
@@ -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"),
docs/reference/node_reference.md
+5
View File
@@ -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. |
docs/reference/pt_train_rnn_model.md
+48
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Training
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptn_rnn.md
+34
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Model
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptn_rnn_linear.md
+36
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Model
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptv_hf_glove_dataset.md
+30
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Training
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptv_sequential_tensor_dataset.md
+22
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Training
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
examples/workflows/rnn_classification.json
+1
View File
File diff suppressed because one or more lines are too long
modules/pytorch_wrapper/hf_glove_dataset.py
+129
View File
@@ -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 "&amp;", "&lt;", etc. to their actual characters
text = unescape(text)
# Remove HTML tags (like <br />, <p>, 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)
modules/pytorch_wrapper/pt_train_model.py
+2 -1
View File
@@ -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()
modules/pytorch_wrapper/pt_train_rnn_model.py
+262
View File
@@ -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))
modules/pytorch_wrapper/ptn_rnn.py
+116
View File
@@ -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,)
modules/pytorch_wrapper/ptn_rnn_linear.py
+184
View File
@@ -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,)
modules/pytorch_wrapper/ptv_hf_glove_dataset.py
+80
View File
@@ -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,)
modules/pytorch_wrapper/ptv_sequential_tensor_dataset.py
+99
View File
@@ -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,)
modules/pytorch_wrapper/rnn_output_reshape.py
+84
View File
@@ -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
modules/pytorch_wrapper/utils.py
+16
View File
@@ -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
pyproject.toml
+1 -1
View File
@@ -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]
requirements.txt
+2
View File
@@ -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
tests/pytorch_wrapper/hf_glove_dataset_test.py
+64
View File
@@ -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()
tests/pytorch_wrapper/pt_train_rnn_linear_model_test.py
+97
View File
@@ -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()
tests/pytorch_wrapper/pt_train_rnn_linear_text_classification_model_test.py
+122
View File
@@ -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()
tests/pytorch_wrapper/pt_train_rnn_model_test.py
+96
View File
@@ -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()
tests/pytorch_wrapper/ptn_rnn_linear_test.py
+91
View File
@@ -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()
tests/pytorch_wrapper/ptn_rnn_test.py
+156
View File
@@ -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()
tests/pytorch_wrapper/ptv_hf_glove_dataset_test.py
+50
View File
@@ -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()
tests/pytorch_wrapper/ptv_sequential_tensor_dataset_test.py
+83
View File
@@ -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()
tests/pytorch_wrapper/rnn_output_reshape_test.py
+152
View File
@@ -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()