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

Added Transformer support

Browse Source
This commit is contained in:
HowToSD
2025-03-28 19:16:55 -07:00
parent e6462b91f8
commit 312f250672
16 changed files with 826 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+6
View File
@@ -110,6 +110,7 @@ from .modules.pytorch_wrapper.pt_to_uint8 import PtToUint8
from .modules.pytorch_wrapper.pt_tokenizer import PtTokenizer
from .modules.pytorch_wrapper.pt_train_classification_model import PtTrainClassificationModel
from .modules.pytorch_wrapper.pt_train_classification_model_lr import PtTrainClassificationModelLr
from .modules.pytorch_wrapper.pt_train_classification_transformer_model import PtTrainClassificationTransformerModel
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
@@ -159,6 +160,7 @@ from .modules.pytorch_wrapper.ptn_conv_model import PtnConvModel
from .modules.pytorch_wrapper.ptn_cross_entropy_loss import PtnCrossEntropyLoss
from .modules.pytorch_wrapper.ptn_embedding import PtnEmbedding
from .modules.pytorch_wrapper.ptn_embedding_rnn_linear import PtnEmbeddingRNNLinear
from .modules.pytorch_wrapper.ptn_embedding_transformer_linear import PtnEmbeddingTransformerLinear
from .modules.pytorch_wrapper.ptn_gru import PtnGRU
from .modules.pytorch_wrapper.ptn_gru_linear import PtnGRULinear
from .modules.pytorch_wrapper.ptn_huber_loss import PtnHuberLoss
@@ -317,6 +319,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
"PtTokenizer": PtTokenizer,
"PtTrainClassificationModel": PtTrainClassificationModel,
"PtTrainClassificationModelLr": PtTrainClassificationModelLr,
"PtTrainClassificationTransformerModel": PtTrainClassificationTransformerModel,
"PtTrainModel": PtTrainModel,
"PtTrainRNNModel": PtTrainRNNModel,
"PtTrainRegressionModel": PtTrainRegressionModel,
@@ -366,6 +369,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
"PtnCrossEntropyLoss": PtnCrossEntropyLoss,
"PtnEmbedding": PtnEmbedding,
"PtnEmbeddingRNNLinear": PtnEmbeddingRNNLinear,
"PtnEmbeddingTransformerLinear": PtnEmbeddingTransformerLinear,
"PtnGRU": PtnGRU,
"PtnGRULinear": PtnGRULinear,
"PtnHuberLoss": PtnHuberLoss,
@@ -524,6 +528,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
"PtTokenizer": "Pt Tokenizer",
"PtTrainClassificationModel": "Pt Train Classification Model",
"PtTrainClassificationModelLr": "Pt Train Classification Model Lr",
"PtTrainClassificationTransformerModel": "Pt Train Classification Transformer Model",
"PtTrainModel": "Pt Train Model",
"PtTrainRNNModel": "Pt Train RNN Model",
"PtTrainRegressionModel": "Pt Train Regression Model",
@@ -573,6 +578,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
"PtnCrossEntropyLoss": "Ptn Cross Entropy Loss",
"PtnEmbedding": "Ptn Embedding",
"PtnEmbeddingRNNLinear": "Ptn Embedding RNN Linear",
"PtnEmbeddingTransformerLinear": "Ptn Embedding Transformer Linear",
"PtnGRU": "Ptn GRU",
"PtnGRULinear": "Ptn GRU Linear",
"PtnHuberLoss": "Ptn Huber Loss",
docs/reference/node_reference.md
+2
View File
@@ -157,6 +157,7 @@
| [Ptn Conv 2d](ptn_conv2d.md) | A convolutional model consisting of a single conv2d layer. |
| [Ptn Conv Model](ptn_conv_model.md) | A convolutional model consisting of multiple convolutional layers. |
| [Ptn Embedding](ptn_embedding.md) | Constructs an embedding layer. |
| [Ptn Embedding Transformer Linear](ptn_embedding_transformer_linear.md) | A Transformer encoder model with a linear head. |
| [Ptn GRU](ptn_gru.md) | A gated recurrent unit (GRU) model consisting of one or more of a recurrent layer. |
| [Ptn GRU Linear](ptn_gru_linear.md) | A recurrent neural network (GRU) model with a linear head. |
| [Ptn Instance Norm 2d](ptn_instance_norm2d.md) | A normalization model to normalize elements over spatial axes within each channel for each sample. |
@@ -234,6 +235,7 @@
| [Pt Predict Regression Model](pt_predict_regression_model.md) | Performs inference on input data. |
| [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 Classification Transformer Model](pt_train_classification_transformer_model.md) | Trains a classification Transformer model using a given dataset, loss function, 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 RNN Model](pt_train_rnn_model.md) | Trains an RNN model using a given dataset, loss function, optimizer, and number of epochs with learning rate decay. |
| [Pt Train Regression Model](pt_train_regression_model.md) | Trains a regression model using a given dataset, optimizer, and number of epochs. |
docs/reference/pt_train_classification_transformer_model.md
+44
View File
@@ -0,0 +1,44 @@
# Pt Train Classification Transformer Model
Trains a classification Transformer 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.
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.
val_loader (PTDATALOADER): DataLoader for validation dataset (optional).
scheduler (PTLRSCHEDULER): Learning rate scheduler.
## Input
| Name | Data type |
|---|---|
| model | Ptmodel |
| 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 |
| 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_embedding_transformer_linear.md
+46
View File
@@ -0,0 +1,46 @@
# Ptn Embedding Transformer Linear
A Transformer encoder model with a linear head.
Args:
num_layers (int): Number of encoder layers in Transformer.
vocabulary_size (int): The vocabulary size of the tokens.
hidden_size (int): Hidden dimension (d_model) for each token of the Transformer.
nhead (int): Number of attention heads.
dim_feedforward (int): Feedforward dimension.
dropout (float): Dropout rate for Transformer layers.
nonlinearity (str): Activation function ('relu' or 'gelu') used in Transformer layers.
layer_norm_eps (float): Epsilon value for layer normalization.
batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature).
norm_first (bool): Whether to apply normalization before attention/feedforward.
bias (bool): Use bias in the Transformer layers.
max_length (int): Maximum input sequence length (for positional embedding).
linear_output_size (int): Output size of the final linear layer.
linear_bias (bool): Use bias in the linear layer.
## Input
| Name | Data type |
|---|---|
| num_layers | Int |
| vocabulary_size | Int |
| hidden_size | Int |
| nhead | Int |
| dim_feedforward | Int |
| dropout | Float |
| nonlinearity | |
| layer_norm_eps | Float |
| batch_first | Boolean |
| norm_first | Boolean |
| bias | Boolean |
| max_length | Int |
| linear_output_size | Int |
| linear_bias | 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_dataset_with_token_encode.md
+2
View File
@@ -7,6 +7,7 @@ A PyTorch Dataset class wrapper for a Hugging Face dataset that converts text to
sample_field_name (str): Field name for text samples.
label_field_name (str): Field name for labels.
encode (PTCALLABLE): The reference to a token encoder function.
remove_html_tags (bool): Remove html tags in text if True.
## Input
| Name | Data type |
@@ -16,6 +17,7 @@ A PyTorch Dataset class wrapper for a Hugging Face dataset that converts text to
| sample_field_name | String |
| label_field_name | String |
| encode | Ptcallable |
| remove_html_tags | Boolean |
## Output
| Data type |
examples/workflows/embedding_transformer_classification.json
+1
View File
File diff suppressed because one or more lines are too long
modules/pytorch_wrapper/hf_dataset_with_token_encode.py
+23 -4
View File
@@ -1,4 +1,6 @@
from typing import Tuple, Callable
from typing import Tuple, Callable, Optional
import re
from html import unescape
import torch
from torch.utils.data import Dataset
from datasets import load_dataset # Hugging Face (not PyTorch)
@@ -15,7 +17,8 @@ class HfDatasetWithTokenEncode(Dataset):
split: str,
sample_field_name: str = "text",
label_field_name: str = "label",
encode: Callable=None) -> None:
encode: Optional[Callable]=None,
remove_html_tags:Optional[bool]=False) -> None:
"""
Initializes the dataset by loading Hugging Face dataset.
@@ -24,12 +27,14 @@ class HfDatasetWithTokenEncode(Dataset):
split (str): The dataset split (e.g., 'train').
sample_field_name (str): Field name for text samples.
label_field_name (str): Field name for labels.
encode (Callable): Token encode function
encode (Optional[Callable]): Token encode function
remove_html_tags (Optional[bool]): Remove html tags in text if True.
"""
self.dataset = load_dataset(dataset_name, split=split)
self.sample_field_name = sample_field_name
self.label_field_name = label_field_name
self.encode = encode
self.remove_html_tags = remove_html_tags
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -46,7 +51,21 @@ class HfDatasetWithTokenEncode(Dataset):
"""
# Get a text sample from database and convert to a word embedding
sample = self.dataset[idx]
tokens, masks = self.encode(sample[self.sample_field_name])
if self.remove_html_tags:
# Decode HTML entities like "&amp;", "&lt;", etc. to their actual characters
text = unescape(sample[self.sample_field_name])
# 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)
# Replace consecutive white space to a single white space
text = re.sub(r" +", " ", text)
else:
text = sample[self.sample_field_name]
tokens, masks = self.encode(text)
# Get a corresponding label
label = torch.tensor(sample[self.label_field_name], dtype=torch.long)
modules/pytorch_wrapper/pt_train_classification_transformer_model.py
+231
View File
@@ -0,0 +1,231 @@
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
class PtTrainClassificationTransformerModel:
"""
Trains a classification Transformer 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.
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.
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", {}),
"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":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,
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,
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 Transformer model.
Args:
model (torch.nn.Module): The PyTorch model to be trained.
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.
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:
if isinstance(x, list):
x = [e.to("cuda") for e in x]
else:
x = x.to("cuda")
y = y.to("cuda")
optimizer.zero_grad()
y_hat = model(x)
if y.dim() == 1:
y = torch.unsqueeze(y, dim=-1)
y = y.to(torch.float32)
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:
if isinstance(x, list):
x = [e.to("cuda") for e in x]
else:
x = x.to("cuda")
y = y.to("cuda")
y_hat = model(x)
if y.dim() == 1:
y = torch.unsqueeze(y, dim=-1)
y = y.to(torch.float32)
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_embedding_transformer_linear.py
+251
View File
@@ -0,0 +1,251 @@
from typing import Any, Dict, Tuple, Union, Callable, Optional
import torch
import torch.nn as nn
import torch
import torch.nn as nn
from typing import Tuple, Union, Callable
class EmbeddingTransformerLinear(nn.Module):
"""
A Transformer encoder-based classification model with a linear output layer.
This module uses a standard PyTorch TransformerEncoder, followed by masked mean pooling
and a linear layer applied to the pooled representation.
Args:
num_layers (int): Number of encoder layers.
vocabulary_size (int): Vocabulary size of the tokens.
hidden_size (int): Hidden size (d_model) of the Transformer.
nhead (int): Number of attention heads.
dim_feedforward (int): Feedforward dimension.
dropout (float): Dropout rate.
nonlinearity (str): Activation function ('relu' or 'gelu').
layer_norm_eps (float): Epsilon for layer normalization.
batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature).
norm_first (bool): Whether normalization comes before sublayers.
bias (bool): Use bias in Transformer layers.
max_length (int): The maximum sequence length used for allocating the positional embedding. This value is not passed to the Transformer itself.
linear_output_size (int): Size of linear layer output.
linear_bias (bool): Use bias in linear layer.
pragma: skip_doc
"""
def __init__(self,
num_encoder_layers: int,
vocabulary_size: int,
d_model: int,
nhead: int,
dim_feedforward: int,
dropout: float,
activation: Union[str, Callable[[torch.Tensor], torch.Tensor]],
layer_norm_eps: float,
batch_first: bool,
norm_first: bool,
bias: bool,
max_length: int,
linear_output_size: int,
linear_bias: bool
):
super().__init__()
self.batch_first = batch_first
self.word_embedding = nn.Embedding(vocabulary_size, d_model)
self.positional_embedding = nn.Embedding(max_length, d_model)
# FIX: Allocate position_ids buffer once to avoid reallocation in each forward pass
# self.register_buffer("position_ids_buffer", torch.arange(max_length)) # (max_length,)
# FIX: use preallocated position_ids_buffer
# position_ids = self.position_ids_buffer[:seq_len].unsqueeze(0).expand(batch_size, -1) # (B, T)
self.position_ids = torch.arange(max_length)
encoder_layer = nn.TransformerEncoderLayer(
d_model=d_model,
nhead=nhead,
dim_feedforward=dim_feedforward,
dropout=dropout,
activation=activation,
layer_norm_eps=layer_norm_eps,
batch_first=batch_first,
norm_first=norm_first,
bias=bias
)
self.encoder = nn.TransformerEncoder(
encoder_layer=encoder_layer,
num_layers=num_encoder_layers
)
self.linear_model = nn.Linear(
in_features=d_model,
out_features=linear_output_size,
bias=linear_bias
)
torch.nn.init.xavier_uniform_(self.linear_model.weight)
def forward(self, inputs: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
"""
Forward pass through the Transformer encoder and linear layer.
Args:
inputs (Tuple[torch.Tensor, torch.Tensor]):
tokens: (B, T) token IDs,
masks: (B, T) attention mask (1 = keep, 0 = pad)
Returns:
torch.Tensor: Output of the linear layer (B, linear_output_size)
"""
tokens, masks = inputs # (B, Tokens), (B, Tokens)
if tokens.dim() != 2:
raise ValueError(f"tokens with rank {tokens.dim()} was passed.")
if masks.dim() != 2:
raise ValueError(f"masks with rank {tokens.dim()} was passed.")
if tokens.size() != masks.size():
raise ValueError(f"tokens size masks size do not match.")
batch_size, seq_len = tokens.size()
token_embeds = self.word_embedding(tokens)
position_ids_repeated = self.position_ids[:seq_len].repeat(batch_size, 1).to(token_embeds.device)
positional_embeds = self.positional_embedding(position_ids_repeated)
if token_embeds.size() != positional_embeds.size():
raise ValueError("token_embeds & positional_embeds size mismatch.")
x = token_embeds + positional_embeds # (B, Seq, Feature)
if not self.batch_first:
x = x.transpose(0, 1) # (Seq, B, Feature)
x = self.encoder(
x,
src_key_padding_mask=~masks.bool() # (B, Seq)
)
if not self.batch_first:
x = x.transpose(0, 1) # (B, Seq, Feature)
# FIX: More memory-efficient masked mean pooling
mask = masks.float() # (B, Seq)
x = (x * mask.unsqueeze(-1)).sum(dim=1) / mask.sum(dim=1, keepdim=True).clamp(min=1e-9) # (B, Feature)
outputs = self.linear_model(x) # (B, linear_output_size)
return outputs
class PtnEmbeddingTransformerLinear:
"""
A Transformer encoder model with a linear head.
Args:
num_layers (int): Number of encoder layers in Transformer.
vocabulary_size (int): The vocabulary size of the tokens.
hidden_size (int): Hidden dimension (d_model) for each token of the Transformer.
nhead (int): Number of attention heads.
dim_feedforward (int): Feedforward dimension.
dropout (float): Dropout rate for Transformer layers.
nonlinearity (str): Activation function ('relu' or 'gelu') used in Transformer layers.
layer_norm_eps (float): Epsilon value for layer normalization.
batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature).
norm_first (bool): Whether to apply normalization before attention/feedforward.
bias (bool): Use bias in the Transformer layers.
max_length (int): Maximum input sequence length (for positional embedding).
linear_output_size (int): Output size of the final linear layer.
linear_bias (bool): Use bias in the 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": {
"num_layers": ("INT", {"default": 6, "min": 1, "max": 256}),
"vocabulary_size": ("INT", {"default": 10000, "min": 1, "max": 1e6}),
"hidden_size": ("INT", {"default": 512, "min": 1, "max": 1e6}),
"nhead": ("INT", {"default": 8, "min": 1, "max": 32}),
"dim_feedforward": ("INT", {"default": 2048, "min": 1, "max": 1e6}),
"dropout": ("FLOAT", {"default": 0.1, "min": 0, "max": 1}),
"nonlinearity": (("gelu", "relu"),),
"layer_norm_eps": ("FLOAT", {"default": 1e-5, "min": 1e-9, "max": 1e-1, "step": 1e-9}),
"batch_first": ("BOOLEAN", {"default": True}),
"norm_first": ("BOOLEAN", {"default": False}),
"bias": ("BOOLEAN", {"default": True}),
"max_length": ("INT", {"default": 512, "min": 1, "max": 1e6}),
"linear_output_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
"linear_bias": ("BOOLEAN", {"default": True}),
}
}
RETURN_TYPES: Tuple[str] = ("PTMODEL",)
FUNCTION: str = "f"
CATEGORY: str = "Training"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self,
num_layers: int,
vocabulary_size: int,
hidden_size: int,
nhead: int,
dim_feedforward: int,
dropout: float,
nonlinearity: str,
layer_norm_eps: float,
batch_first: bool,
norm_first: bool,
bias: bool,
max_length: int,
linear_output_size: int,
linear_bias: bool
) -> Tuple[nn.Module]:
"""
Constructs a Transformer model with a linear output layer.
Args:
num_layers (int): Number of encoder layers.
vocabulary_size (int): Vocabulary size of the tokens.
hidden_size (int): Hidden size (d_model) of the Transformer.
nhead (int): Number of attention heads.
dim_feedforward (int): Feedforward dimension.
dropout (float): Dropout rate.
nonlinearity (str): Activation function ('relu' or 'gelu').
layer_norm_eps (float): Epsilon for layer normalization.
batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature).
norm_first (bool): Whether normalization comes before sublayers.
bias (bool): Use bias in Transformer layers.
max_length (int): Maximum sequence length.
linear_output_size (int): Size of linear layer output.
linear_bias (bool): Use bias in linear layer.
Returns:
Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
"""
with torch.inference_mode(False):
model = EmbeddingTransformerLinear(
num_encoder_layers=num_layers,
vocabulary_size=vocabulary_size,
d_model=hidden_size,
nhead=nhead,
dim_feedforward=dim_feedforward,
dropout=dropout,
activation=nonlinearity,
layer_norm_eps=layer_norm_eps,
batch_first=batch_first,
norm_first=norm_first,
bias=bias,
max_length=max_length,
linear_output_size=linear_output_size,
linear_bias=linear_bias
)
return (model,)
modules/pytorch_wrapper/pto_lr_scheduler_cosine_annealing.py
+1 -1
View File
@@ -54,7 +54,7 @@ class PtoLrSchedulerCosineAnnealing:
"required": {
"optimizer": ("PTOPTIMIZER", {}),
"num_epochs": ("INT", {"default": 10, "min": 1, "max": int(1e8)}),
"minimum_lr": ("FLOAT", {"default": 0, "min": 0, "max": 1.0, "step": 1e-8}),
"minimum_lr": ("FLOAT", {"default": 0, "min": 0, "max": 1.0, "step": 1e-9}),
}
}
modules/pytorch_wrapper/ptv_hf_dataset_with_token_encode.py
+7 -2
View File
@@ -12,6 +12,7 @@ class PtvHfDatasetWithTokenEncode:
sample_field_name (str): Field name for text samples.
label_field_name (str): Field name for labels.
encode (PTCALLABLE): The reference to a token encoder function.
remove_html_tags (bool): Remove html tags in text if True.
category: PyTorch wrapper - Training
"""
@@ -31,6 +32,7 @@ class PtvHfDatasetWithTokenEncode:
"sample_field_name": ("STRING", {"default": "text","multiline": False}),
"label_field_name": ("STRING", {"default": "label","multiline": False}),
"encode": ("PTCALLABLE", {}),
"remove_html_tags": ("BOOLEAN", {"default": False}),
}
}
@@ -47,7 +49,8 @@ class PtvHfDatasetWithTokenEncode:
split: str,
sample_field_name: str,
label_field_name: str,
encode: Callable) -> Tuple:
encode: Callable,
remove_html_tags:bool) -> Tuple:
"""
Loads a dataset from Hugging Face with specified parameters.
@@ -57,6 +60,7 @@ class PtvHfDatasetWithTokenEncode:
sample_field_name (str): Field name for text samples.
label_field_name (str): Field name for labels.
encode (Callable): The reference to a token encoder function.
remove_html_tags (bool): Remove html tags in text if True.
Returns:
Tuple: A tuple containing the dataset instance.
@@ -71,5 +75,6 @@ class PtvHfDatasetWithTokenEncode:
split=split,
sample_field_name=sample_field_name,
label_field_name=label_field_name,
encode=encode)
encode=encode,
remove_html_tags=remove_html_tags)
return (dc,)
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.3.4"
version = "3.4.0"
license = {file = "LICENSE"}
[project.urls]
tests/pytorch_wrapper/hf_dataset_with_tokenize_test.py
+27
View File
@@ -37,6 +37,13 @@ class TestHfDatasetWithTokenEncode(unittest.TestCase):
"train", #split
encode=encode,
)
self.ds_remove_html = HfDatasetWithTokenEncode(
"imdb", # dataset_name
"train", #split
encode=encode,
remove_html_tags=True
)
def test_loading(self):
it = iter(self.ds)
@@ -58,5 +65,25 @@ class TestHfDatasetWithTokenEncode(unittest.TestCase):
self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.")
def test_loading_remove_html(self):
it = iter(self.ds_remove_html)
(sample, mask), label = next(it)
# Compare text
actual = sample.size()
expected = torch.Size((self.max_length, ))
self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.")
# Compare mask
actual = mask.size()
expected = torch.Size((self.max_length, ))
self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.")
# Compare label
actual = label
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_embedding_transformer_linear_text_classification_model_test.py
+139
View File
@@ -0,0 +1,139 @@
"""
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_embedding_transformer_linear import PtnEmbeddingTransformerLinear
from pytorch_wrapper.hf_dataset_with_token_encode import HfDatasetWithTokenEncode
from sentencepiece_wrapper.sp_load_model import SpLoadModel
from sentencepiece_wrapper.sp_encode import SpEncode
from pytorch_wrapper.pto_adamw import PtoAdamW
from pytorch_wrapper.pt_data_loader import PtDataLoader
from pytorch_wrapper.pt_train_classification_transformer_model import PtTrainClassificationTransformerModel
from pytorch_wrapper.pt_save_model import PtSaveModel
from pytorch_wrapper.pt_load_model import PtLoadModel
from pytorch_wrapper.pto_lr_scheduler_cosine_annealing import PtoLrSchedulerCosineAnnealing
from pytorch_wrapper.ptn_bce_with_logits_loss import PtnBCEWithLogitsLoss
from pytorch_wrapper.utils import set_seed
class TestPtTrainEmbeddingTransformerLinearTextClassificationModel(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]
self.load_node = SpLoadModel()
self.sp = self.load_node.f("spiece.model")[0]
self.encode_node = SpEncode()
self.max_length=512
encode = self.encode_node.f(
spmodel=self.sp,
padding=True,
padding_method="max_length",
padding_value=0,
truncation=True,
max_length=self.max_length
)[0]
train_dataset = HfDatasetWithTokenEncode(
"imdb", # dataset_name
"train", #split
encode=encode,
remove_html_tags=True
)
test_dataset = HfDatasetWithTokenEncode(
"imdb", # dataset_name
"test", #split
encode=encode,
remove_html_tags=True
)
train_data_loader_node = PtDataLoader()
test_data_loader_node = PtDataLoader()
self.train_loader = train_data_loader_node.f(
dataset=train_dataset,
batch_size=32,
shuffle=True,
parameters='{"num_workers":1}'
)[0]
self.test_loader = test_data_loader_node.f(
dataset=test_dataset,
batch_size=32,
shuffle=False,
parameters='{"num_workers":1}'
)[0]
self.model_path = "transformer_text_classification_test.pt"
self.model_node = PtnEmbeddingTransformerLinear()
self.model = self.model_node.f(
2, # num_encoder
32000, # vocabulary_size
256, # hidden_size
8, # nhead
512, # dim_feedforward
0.5, # dropout
"gelu", # activation
1e-5, # layer_norm_eps,
True, # batch_first
False, # norm_first
True, # bias
self.max_length, # max_length of seq
1, # linear_output_size
True # linear_bias
)[0]
self.optimizer = PtoAdamW().f(self.model, 3e-4, 0.9, 0.999,
0.01, # weight decay
False # amsgrad
)[0]
self.minimum_lr = 1e-9
self.epochs = 10
self.scheduler = PtoLrSchedulerCosineAnnealing().f(self.optimizer,
self.epochs,
self.minimum_lr)[0]
self.trainer = PtTrainClassificationTransformerModel()
self.save_model = PtSaveModel()
self.load_model = PtLoadModel()
self.linear_head = True
@unittest.skipIf(os.getenv("RUN_SKIPPED_TESTS") != "1", "Skipping unless explicitly enabled")
def test_1(self):
trained_model = self.trainer.f(
self.model,
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,
val_loader=self.test_loader)[0]
self.save_model.f(trained_model, self.model_path)
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/ptn_embedding_transformer_linear_test.py
+43
View File
@@ -0,0 +1,43 @@
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_embedding_transformer_linear import PtnEmbeddingTransformerLinear
class TestPtnEmbeddingTransformerLinear(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.model_node = PtnEmbeddingTransformerLinear()
def test_1(self):
self.model = self.model_node.f(
1, # num_encoder
32000, # vocabulary_size
32, # hidden_size
8, # nhead
64, # dim_feedforward
0.1, # dropout
"gelu", # activation
1e-5, # layer_norm_eps,
True, # batch_first
False, # norm_first
True, # bias
512, # max_length of seq
1, # linear_output_size
True # linear_bias
)[0]
x = torch.randint(0, 32000, (8 * 77,)).to(torch.int64)
x = x.view(8, 77)
masks = torch.ones_like(x)
outputs = self.model((x, masks))
self.assertEqual(outputs.size(), torch.Size([8, 1]))
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/ptv_hf_dataset_with_token_encode_test.py
+2 -1
View File
@@ -36,7 +36,8 @@ class TestPtvHfDatasetWithTokenEncode(unittest.TestCase):
"train", #split
"text",
"label",
encode
encode,
False # remove_html_tags
)[0]