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

Added GRU support.

Browse Source
This commit is contained in:
HowToSD
2025-03-25 19:25:13 -07:00
parent d9a2b827a5
commit e6abdaebdc
19 changed files with 722 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/pytorch_wrapper/hf_glove_dataset.py
+3
View File
@@ -72,6 +72,9 @@ class HfGloveDataset(Dataset):
# This strips punctuation like "!", ".", "(", but keeps words and numbers
text = re.sub(r"[^a-zA-Z0-9\s]", " ", text)
# Replace consecutive white space to a single white space
text = re.sub(r" +", " ", text)
# Convert to lowercase and split by whitespace to get token list
return text.lower().split()
modules/pytorch_wrapper/pt_train_rnn_model.py
+14 -3
View File
@@ -9,6 +9,7 @@ from .rnn_output_reshape import extract_rnn_return_value_with_adjusted_label
class PtTrainRNNModel:
"""
Pt Train RNN Model:
Trains an RNN model using a given dataset, loss function, optimizer, and number of epochs with learning rate decay.
Args:
@@ -26,6 +27,7 @@ class PtTrainRNNModel:
use_valid_token_mean (bool): Computes the mean of the RNN output for non-zero inputs and feeds to the linear layer.
val_loader (PTDATALOADER): DataLoader for validation dataset (optional).
scheduler (PTLRSCHEDULER): Learning rate scheduler.
h_0 (TENSOR): Optional initial hidden state. For unidirectional models, the shape is [num_layers, batch_size, hidden_size]. For bidirectional models, it is [2 * num_layers, batch_size, hidden_size]. If not provided, a zero tensor is used.
category: PyTorch wrapper - Training
"""
@@ -56,6 +58,7 @@ class PtTrainRNNModel:
"optional": {
"scheduler": ("PTLRSCHEDULER", {}),
"val_loader": ("PTDATALOADER", {}),
"h_0": ("TENSOR", {}),
}
}
@@ -82,7 +85,8 @@ class PtTrainRNNModel:
classification_metrics:bool,
use_valid_token_mean: bool,
scheduler: Optional[Union[torch.optim.lr_scheduler._LRScheduler, torch.optim.lr_scheduler.ReduceLROnPlateau]] = None,
val_loader: Optional[torch.utils.data.DataLoader] = None) -> Tuple[nn.Module, torch.Tensor, torch.Tensor]:
val_loader: Optional[torch.utils.data.DataLoader] = None,
h_0: Optional[torch.Tensor] = None) -> Tuple[nn.Module, torch.Tensor, torch.Tensor]:
"""
Trains an RNN model.
@@ -101,6 +105,7 @@ class PtTrainRNNModel:
use_valid_token_mean (bool): Computes the mean of the RNN output for non-zero inputs and feeds to the linear layer.
val_loader (Optional[torch.utils.data.DataLoader]): DataLoader for validation dataset (optional).
scheduler (Optional[Union[torch.optim.lr_scheduler._LRScheduler, torch.optim.lr_scheduler.ReduceLROnPlateau]]): Learning rate scheduler.
h_0 (optional[torch.Tensor]): Optional initial hidden state. For unidirectional models, the shape is [num_layers, batch_size, hidden_size]. For bidirectional models, it is [2 * num_layers, batch_size, hidden_size]. If not provided, a zero tensor is used.
Returns:
Tuple[torch.nn.Module, torch.Tensor, torch.Tensor]: A tuple containing the trained model,
@@ -142,14 +147,20 @@ class PtTrainRNNModel:
y = y.to("cuda")
optimizer.zero_grad()
if linear_head:
y_hat = model(x)
if h_0:
y_hat = model(x, hx=h_0) # Note that both RNN and GRU use hx. See PyTorch code. https://github.com/pytorch/pytorch/blob/v2.6.0/torch/nn/modules/rnn.py
else:
y_hat = model(x)
if y.dim() == 3:
y = y[:,-1,:] # Grab last token
elif y.dim() == 1:
y = torch.unsqueeze(y, dim=-1)
y = y.to(torch.float32)
else:
output_seq, _ = model(x)
if h_0:
output_seq, _ = model(x, hx=h_0)
else:
output_seq, _ = model(x)
outputs, adjusted_labels = extract_rnn_return_value_with_adjusted_label(
x,
y,
modules/pytorch_wrapper/ptn_gru.py
+88
View File
@@ -0,0 +1,88 @@
from typing import Any, Dict, Tuple
import torch
import torch.nn as nn
class PtnGRU:
"""
Ptn GRU:
A gated recurrent unit (GRU) model consisting of one or more of a recurrent layer.
Args:
input_size (int): The number of input features.
hidden_size (int): The number of output features of the hidden layer matrix.
num_layers (int): Number of hidden layers.
bias (bool): Use bias or not.
batch_first (bool): Indicates the shape of the input. Input is assumed to be a rank-3 tensor. If True, the shape is [Batch, Seq, Token]. If False, the shape is [Seq, Batch, Token]. Note that default is set to `True` unlike PyTorch's GRU default model parameter.
dropout (float): Dropout probability to be applied to the output of intermediate layers if non-zero. This will not be applied to the last layer's output.
bidirectional (bool): Processes the input in a backward direction and append the result to the output of the forward direction.
category: PyTorch wrapper - Model
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"input_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
"hidden_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
"num_layers": ("INT", {"default": 1, "min": 1, "max": 1e3}),
"bias": ("BOOLEAN", {"default": True}),
"batch_first": ("BOOLEAN", {"default": True}),
"dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1}),
"bidirectional": ("BOOLEAN", {"default": False}),
}
}
RETURN_TYPES: Tuple[str] = ("PTMODEL",)
FUNCTION: str = "f"
CATEGORY: str = "Training"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self,
input_size: int,
hidden_size: int,
num_layers: int,
bias: bool,
batch_first: bool,
dropout: float,
bidirectional: bool
) -> Tuple[nn.Module]:
"""
Constructs a GRU model.
Args:
input_size (int): The number of input features.
hidden_size (int): The number of output features of the hidden layer matrix.
num_layers (int): Number of hidden layers.
bias (bool): Use bias or not.
batch_first (bool): Indicates the shape of the input. Input is assumed to be a rank-3 tensor. If True, the shape is [Batch, Seq, Token]. If False, the shape is [Seq, Batch, Token].
dropout (float): Dropout probability to be applied to the output of intermediate layers if non-zero. This will not be applied to the last layer's output.
bidirectional (bool): Processes the input in a backward direction and append the result to the output of the forward direction.
Returns:
Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
"""
with torch.inference_mode(False):
model = nn.GRU(
input_size=input_size,
hidden_size=hidden_size,
num_layers=num_layers,
bias=bias,
batch_first=batch_first,
dropout=dropout,
bidirectional=bidirectional
)
# TODO: Xaviar normal weight initialization is not
# set up for GRU. If performance turning is needed, consider adding it here.
return (model,)
modules/pytorch_wrapper/ptn_gru_linear.py
+163
View File
@@ -0,0 +1,163 @@
from typing import Any, Dict, Tuple
import torch
import torch.nn as nn
from pytorch_wrapper.rnn_output_reshape import extract_rnn_return_value_with_adjusted_label
class GRULinear(nn.Module):
"""
A recurrent neural network (GRU) with a linear output layer.
This module uses a standard PyTorch GRU, followed by a linear layer applied to the final output tensor.
pragma: skip_doc
"""
def __init__(self,
input_size: int,
hidden_size: int,
num_layers: int,
bias: bool,
batch_first: bool,
dropout: float,
bidirectional: bool,
linear_output_size: int,
linear_bias: bool
):
super().__init__()
self.gru_model = nn.GRU(
input_size=input_size,
hidden_size=hidden_size,
num_layers=num_layers,
bias=bias,
batch_first=batch_first,
dropout=dropout,
bidirectional=bidirectional
)
self.bidirectional = bidirectional
self.batch_first = batch_first
directions = 2 if bidirectional else 1
self.linear_model = nn.Linear(
in_features=hidden_size * directions,
out_features=linear_output_size,
bias=linear_bias
)
torch.nn.init.xavier_uniform_(self.linear_model.weight)
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
"""
Forward pass through the GRU and the linear layer.
Args:
inputs (torch.Tensor): Input tensor of shape (batch, seq, input_size) if batch_first=True.
Returns:
torch.Tensor: Output of the linear layer.
"""
out_seq, _ = self.gru_model(inputs)
x = extract_rnn_return_value_with_adjusted_label(
inputs, # This is needed to find out padded tokens
None,
out_seq,
self.bidirectional,
self.batch_first,
False,
return_valid_token_mean=True
)
outputs = self.linear_model(x)
return outputs
class PtnGRULinear:
"""
Ptn GRU Linear:
A recurrent neural network (GRU) model with a linear head.
Args:
input_size (int): The number of input features.
hidden_size (int): The number of output features of the hidden layer matrix.
num_layers (int): Number of hidden layers in GRU.
bias (bool): Use bias or not.
batch_first (bool): Indicates the shape of the input. Input is assumed to be a rank-3 tensor. If True, the shape is [Batch, Seq, Token]. If False, the shape is [Seq, Batch, Token]. Note that default is set to `True` unlike PyTorch's GRU default model parameter.
dropout (float): Dropout probability to be applied to the output of intermediate layers if non-zero. This will not be applied to the last layer's output.
bidirectional (bool): Processes the input in a backward direction and append the result to the output of the forward direction.
linear_output_size (int): The number of output features of the linear layer.
linar_bias (bool): Use bias or not in linear layer.
category: PyTorch wrapper - Model
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"input_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
"hidden_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
"num_layers": ("INT", {"default": 1, "min": 1, "max": 1e3}),
"bias": ("BOOLEAN", {"default": True}),
"batch_first": ("BOOLEAN", {"default": True}),
"dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1}),
"bidirectional": ("BOOLEAN", {"default": False}),
"linear_output_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
"linear_bias": ("BOOLEAN", {"default": True}),
}
}
RETURN_TYPES: Tuple[str] = ("PTMODEL",)
FUNCTION: str = "f"
CATEGORY: str = "Training"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self,
input_size: int,
hidden_size: int,
num_layers: int,
bias: bool,
batch_first: bool,
dropout: float,
bidirectional: bool,
linear_output_size: int,
linear_bias: bool
) -> Tuple[nn.Module]:
"""
Constructs an GRU model.
Args:
input_size (int): The number of input features.
hidden_size (int): The number of output features of the hidden layer matrix.
num_layers (int): Number of hidden layers.
bias (bool): Use bias or not.
batch_first (bool): Indicates the shape of the input. Input is assumed to be a rank-3 tensor. If True, the shape is [Batch, Seq, Token]. If False, the shape is [Seq, Batch, Token].
dropout (float): Dropout probability to be applied to the output of intermediate layers if non-zero. This will not be applied to the last layer's output.
bidirectional (bool): Processes the input in a backward direction and append the result to the output of the forward direction.
linear_output_size (int): The number of output features of the linear layer.
linar_bias (bool): Use bias or not in linear layer.
Returns:
Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
"""
with torch.inference_mode(False):
model = GRULinear(
input_size=input_size,
hidden_size=hidden_size,
num_layers=num_layers,
bias=bias,
batch_first=batch_first,
dropout=dropout,
bidirectional=bidirectional,
linear_output_size=linear_output_size,
linear_bias=linear_bias
)
return (model,)
modules/pytorch_wrapper/ptn_rnn.py
-7
View File
@@ -105,12 +105,5 @@ class PtnRNN:
torch.nn.init.xavier_normal_(weight_ih)
torch.nn.init.xavier_normal_(weight_hh)
# bias_ih = getattr(model, f"bias_ih_l{layer}{suffix}")
# bias_hh = getattr(model, f"bias_hh_l{layer}{suffix}")
# torch.nn.init.uniform_(weight_ih, -0.05, 0.05)
# torch.nn.init.uniform_(weight_hh, -0.05, 0.05)
# nn.init.zeros_(bias_ih)
# nn.init.zeros_(bias_hh)
return (model,)