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

Added MHA and related nodes.

Browse Source
This commit is contained in:
HowToSD
2025-04-06 18:18:44 -07:00
parent a2490be96f
commit ca1644d2f5
29 changed files with 1219 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+18
View File
@@ -157,6 +157,7 @@ from .modules.pytorch_wrapper.ptn_batch_norm2d import PtnBatchNorm2d
from .modules.pytorch_wrapper.ptn_bce_loss import PtnBCELoss
from .modules.pytorch_wrapper.ptn_bce_with_logits_loss import PtnBCEWithLogitsLoss
from .modules.pytorch_wrapper.ptn_chained_model import PtnChainedModel
from .modules.pytorch_wrapper.ptn_chained_model_with_attention_mask import PtnChainedModelWithAttentionMask
from .modules.pytorch_wrapper.ptn_conv2d import PtnConv2d
from .modules.pytorch_wrapper.ptn_conv_model import PtnConvModel
from .modules.pytorch_wrapper.ptn_cross_entropy_loss import PtnCrossEntropyLoss
@@ -176,11 +177,16 @@ from .modules.pytorch_wrapper.ptn_linear import PtnLinear
from .modules.pytorch_wrapper.ptn_linear_model import PtnLinearModel
from .modules.pytorch_wrapper.ptn_lstm import PtnLSTM
from .modules.pytorch_wrapper.ptn_lstm_linear import PtnLSTMLinear
from .modules.pytorch_wrapper.ptn_masked_mean_pooling import PtnMaskedMeanPooling
from .modules.pytorch_wrapper.ptn_max_pool2d import PtnMaxPool2d
from .modules.pytorch_wrapper.ptn_model_with_closure import PtnModelWithClosure
from .modules.pytorch_wrapper.ptn_mse_loss import PtnMSELoss
from .modules.pytorch_wrapper.ptn_multihead_attention import PtnMultiheadAttention
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_residual_connection_model import PtnResidualConnectionModel
from .modules.pytorch_wrapper.ptn_residual_connection_model_with_attention_mask import PtnResidualConnectionModelWithAttentionMask
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
@@ -371,6 +377,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
"PtnBCEWithLogitsLoss": PtnBCEWithLogitsLoss,
"PtnBatchNorm2d": PtnBatchNorm2d,
"PtnChainedModel": PtnChainedModel,
"PtnChainedModelWithAttentionMask": PtnChainedModelWithAttentionMask,
"PtnConv2d": PtnConv2d,
"PtnConvModel": PtnConvModel,
"PtnCrossEntropyLoss": PtnCrossEntropyLoss,
@@ -391,12 +398,17 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
"PtnLinear": PtnLinear,
"PtnLinearModel": PtnLinearModel,
"PtnMSELoss": PtnMSELoss,
"PtnMaskedMeanPooling": PtnMaskedMeanPooling,
"PtnMaxPool2d": PtnMaxPool2d,
"PtnModelWithClosure": PtnModelWithClosure,
"PtnMultiheadAttention": PtnMultiheadAttention,
"PtnNLLLoss": PtnNLLLoss,
"PtnPreAddChannelAxis": PtnPreAddChannelAxis,
"PtnPreFlatten": PtnPreFlatten,
"PtnRNN": PtnRNN,
"PtnRNNLinear": PtnRNNLinear,
"PtnResidualConnectionModel": PtnResidualConnectionModel,
"PtnResidualConnectionModelWithAttentionMask": PtnResidualConnectionModelWithAttentionMask,
"PtnResnetModel": PtnResnetModel,
"PtnSmoothL1Loss": PtnSmoothL1Loss,
"PtoAdam": PtoAdam,
@@ -585,6 +597,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
"PtnBCEWithLogitsLoss": "Ptn BCE With Logits Loss",
"PtnBatchNorm2d": "Ptn Batch Norm 2d",
"PtnChainedModel": "Ptn Chained Model",
"PtnChainedModelWithAttentionMask": "Ptn Chained Model With Attention Mask",
"PtnConv2d": "Ptn Conv 2d",
"PtnConvModel": "Ptn Conv Model",
"PtnCrossEntropyLoss": "Ptn Cross Entropy Loss",
@@ -605,12 +618,17 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
"PtnLinear": "Ptn Linear",
"PtnLinearModel": "Ptn Linear Model",
"PtnMSELoss": "Ptn MSE Loss",
"PtnMaskedMeanPooling": "Ptn Masked Mean Pooling",
"PtnMaxPool2d": "Ptn Max Pool 2d",
"PtnModelWithClosure": "Ptn Model With Closure",
"PtnMultiheadAttention": "Ptn Multihead Attention",
"PtnNLLLoss": "Ptn NLL Loss",
"PtnPreAddChannelAxis": "Ptn Pre Add Channel Axis",
"PtnPreFlatten": "Ptn Pre Flatten",
"PtnRNN": "Ptn RNN",
"PtnRNNLinear": "Ptn RNN Linear",
"PtnResidualConnectionModel": "Ptn Residual Connection Model",
"PtnResidualConnectionModelWithAttentionMask": "Ptn Residual Connection Model With Attention Mask",
"PtnResnetModel": "Ptn Resnet Model",
"PtnSmoothL1Loss": "Ptn Smooth L1 Loss",
"PtoAdam": "Pto Adam",
dev/scripts/reference_generate.py
+1
View File
@@ -115,6 +115,7 @@ def process_python_file(file_path: str, output_dir: str, node_list: Dict[str, Li
category = class_info["category"]
file_name = os.path.splitext(os.path.basename(file_path))[0] + ".md"
file_name = file_name.replace(" ", "") # Remove white space
output_path = os.path.join(output_dir, file_name)
with open(output_path, "w", encoding="utf-8") as md_file:
docs/reference/node_reference.md
+9
View File
@@ -8,6 +8,10 @@
| Node | Description |
| --- | --- |
| [Pt Save Model](pt_save_model.md) | A wrapper class for saving a PyTorch model. |
## Miscellaneous
| Node | Description |
| --- | --- |
| [Ptn Model With Closure](ptn_model_with_closure.md) | A model that is followed by a closure function. |
## PyTorch wrapper - Arithmetic operations
| Node | Description |
| --- | --- |
@@ -154,6 +158,7 @@
| [Ptn Avg Pool 2d](ptn_avg_pool2d.md) | An avgpool layer. |
| [Ptn Batch Norm 2d](ptn_batch_norm2d.md) | A normalization model to normalize over the batch and spatial axes for each channel. |
| [Ptn Chained Model](ptn_chained_model.md) | Constructs a chained PyTorch model. |
| [Ptn Chained Model With Attention Mask](ptn_chained_model_with_attention_mask.md) | A chained model that sequentially applies `model_a` and `model_b`, |
| [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. |
@@ -168,12 +173,16 @@
| [Ptn Layer Norm](ptn_layer_norm.md) | A normalization model to normalize elements over specified axes. |
| [Ptn Linear](ptn_linear.md) | A linear model consisting of a single dense layer. |
| [Ptn Linear Model](ptn_linear_model.md) | A linear model consisting of dense layers. |
| [Ptn Masked Mean Pooling](ptn_masked_mean_pooling.md) | Constructs a masked mean pooling layer. |
| [Ptn Max Pool 2d](ptn_max_pool2d.md) | A maxpool layer. |
| [Ptn Multihead Attention](ptn_multihead_attention.md) | A Multihead attention model. |
| [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_embedding_rnn_linear.md) | A recurrent neural network (RNN) model with a linear head. |
| [Ptn RNN Linear](ptn_rnn_linear.md) | A recurrent neural network (RNN) model with a linear head. |
| [Ptn Residual Connection Model](ptn_residual_connection_model.md) | A model that saves the input and add to the output of the specified model. |
| [Ptn Residual Connection Model With Attention Mask](ptn_residual_connection_model_with_attention_mask.md) | A model that saves the input and add to the output of the specified model. |
| [Ptn Resnet Model](ptn_resnet_model.md) | A Resnet model consisting of multiple Resnet layers. |
## PyTorch wrapper - Optimizer
| Node | Description |
docs/reference/pt_train_classification_transformer_model.md
+4 -1
View File
@@ -13,7 +13,9 @@ Trains a classification Transformer model using a given dataset, loss function,
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.
scheduler (PTLRSCHEDULER): Learning rate scheduler.
unpack_input (Optional[bool]): If True and the input is a list, unpack it when calling the model.
For example, if the dataloader returns [a, b] and your model's forward is defined as forward(a, b), set this to True.
## Input
| Name | Data type |
@@ -30,6 +32,7 @@ Trains a classification Transformer model using a given dataset, loss function,
| classification_metrics | Boolean |
| scheduler | Ptlrscheduler |
| val_loader | Ptdataloader |
| unpack_input | Boolean |
## Output
| Data type |
docs/reference/ptn_chained_model_with_attention_mask.md
+29
View File
@@ -0,0 +1,29 @@
# Ptn Chained Model With Attention Mask
A chained model that sequentially applies `model_a` and `model_b`,
optionally followed by a closure function.
### Fields:
- `model_a`: First model in the chain.
- `model_b`: Second model in the chain.
- `model_a_mask_req`: True if model_a requires attention_mask for input
- `model_b_mask_req`: True if model_b requires attention_mask for input
- closure: Optional differentiable function to be called at the end of forward.
## Input
| Name | Data type |
|---|---|
| model_a | Ptmodel |
| model_b | Ptmodel |
| model_a_mask_req | Boolean |
| model_b_mask_req | Boolean |
| closure | Ptcallable |
## 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_embedding.md
+1 -1
View File
@@ -10,7 +10,7 @@ Constructs an embedding layer.
| Name | Data type |
|---|---|
| num_embeddings | Int |
| embedding_size | Int |
| embedding_dim | Int |
| padding_idx | Int |
## Output
docs/reference/ptn_masked_mean_pooling.md
+13
View File
@@ -0,0 +1,13 @@
# Ptn Masked Mean Pooling
Constructs a masked mean pooling layer.
Computes the mean of the input excluding the padding values using the given mask.
## 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_model_with_closure.md
+22
View File
@@ -0,0 +1,22 @@
# Ptn Model With Closure
A model that is followed by a closure function.
### Fields:
- `model`: The model in the chain.
- `closure`: Differentiable function to be called at the end of forward. category: PyTorch wrapper - Model
## Input
| Name | Data type |
|---|---|
| model | Ptmodel |
| closure | Ptcallable |
## Output
| Data type |
|---|
| Ptmodel |
<HR>
Category: Miscellaneous
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptn_multihead_attention.md
+36
View File
@@ -0,0 +1,36 @@
# Ptn Multihead Attention
A Multihead attention model.
Args:
embed_dim (int): Hidden dimension (d_model) for each token of the Transformer.
num_heads (int): Number of attention heads.
dropout (float): Dropout rate for output.
bias (bool): If True, bias will be added to the linear projections of query, key, value, and the final output.
add_bias_kv (bool): If True, bias will be added to k and v sequences.
add_zero_attn (bool): If True, a zero batch will be added to k and v sequences.
kdim (int): Dimension of k.
vdim (int): Dimension of v.
batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature).
## Input
| Name | Data type |
|---|---|
| embed_dim | Int |
| num_heads | Int |
| dropout | Float |
| bias | Boolean |
| add_bias_kv | Boolean |
| add_zero_attn | Boolean |
| kdim | Int |
| vdim | Int |
| batch_first | 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_residual_connection_model.md
+23
View File
@@ -0,0 +1,23 @@
# Ptn Residual Connection Model
A model that saves the input and add to the output of the specified model.
optionally followed by a closure function.
### Fields:
- `model`: A model to process input.
- `closure`: Optional differentiable function to be called at the end of forward.
## Input
| Name | Data type |
|---|---|
| model | Ptmodel |
| closure | Ptcallable |
## 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_residual_connection_model_with_attention_mask.md
+23
View File
@@ -0,0 +1,23 @@
# Ptn Residual Connection Model With Attention Mask
A model that saves the input and add to the output of the specified model.
optionally followed by a closure function.
### Fields:
- `model`: A model to process input.
- `closure`: Optional differentiable function to be called at the end of forward.
## Input
| Name | Data type |
|---|---|
| model | Ptmodel |
| closure | Ptcallable |
## 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_local_dataset.md
+2 -1
View File
@@ -5,7 +5,8 @@ A PyTorch Dataset class wrapper for a Hugging Face dataset to load a dataset tha
file_path (str): Path to the dataset file. Can be an absolute path or a relative path.
If a relative path is provided, it is resolved from the `datasets` directory under this extension.
For example, to load `datasets/foo/bar_train.jsonp`, specify `foo/bar_train.jsonp` or the absolute path.
file_format (str): File format (e.g. json if the file is in jsonl format). sample_field_name (str): Field name for text samples.
file_format (str): File format (e.g. json if the file is in jsonl format).
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.
examples/workflows/mha.json
+1
View File
File diff suppressed because one or more lines are too long
modules/pytorch_wrapper/pt_train_classification_transformer_model.py
+16 -5
View File
@@ -24,6 +24,9 @@ class PtTrainClassificationTransformerModel:
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.
unpack_input (Optional[bool]): If True and the input is a list, unpack it when calling the model.
For example, if the dataloader returns [a, b] and your model's forward is defined as forward(a, b), set this to True.
category: PyTorch wrapper - Training
"""
@@ -51,7 +54,8 @@ class PtTrainClassificationTransformerModel:
},
"optional": {
"scheduler": ("PTLRSCHEDULER", {}),
"val_loader": ("PTDATALOADER", {})
"val_loader": ("PTDATALOADER", {}),
"unpack_input": ("BOOLEAN", {"default":False})
}
}
@@ -76,7 +80,8 @@ class PtTrainClassificationTransformerModel:
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]:
val_loader: Optional[torch.utils.data.DataLoader] = None,
unpack_input: Optional[bool] = False) -> Tuple[nn.Module, torch.Tensor, torch.Tensor]:
"""
Trains an Transformer model.
@@ -93,7 +98,7 @@ class PtTrainClassificationTransformerModel:
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.
unpack_input (Optional[bool]): Unpack input (x) if input is a list.
Returns:
Tuple[torch.nn.Module, torch.Tensor, torch.Tensor]: A tuple containing the trained model,
training loss history, and validation loss history.
@@ -135,7 +140,10 @@ class PtTrainClassificationTransformerModel:
y = y.to("cuda")
optimizer.zero_grad()
y_hat = model(x)
if unpack_input:
y_hat = model(*x)
else:
y_hat = model(x)
if y.dim() == 1:
y = torch.unsqueeze(y, dim=-1)
y = y.to(torch.float32)
@@ -167,7 +175,10 @@ class PtTrainClassificationTransformerModel:
x = x.to("cuda")
y = y.to("cuda")
y_hat = model(x)
if unpack_input:
y_hat = model(*x)
else:
y_hat = model(x)
if y.dim() == 1:
y = torch.unsqueeze(y, dim=-1)
y = y.to(torch.float32)
modules/pytorch_wrapper/ptn_chained_model.py
+1 -1
View File
@@ -44,7 +44,7 @@ class PtnChainedModelDef(nn.Module):
torch.Tensor: Output tensor after passing through `model_a`, `model_b`, and optionally `closure`.
"""
x = self.model_a(inputs)
x = self.model_b(x) # Fixed the incorrect use of `inputs` instead of `x`
x = self.model_b(x)
if self.closure:
x = self.closure(x)
return x
modules/pytorch_wrapper/ptn_chained_model_with_attention_mask.py
+128
View File
@@ -0,0 +1,128 @@
from typing import Any, Dict, Tuple, Callable, Optional
import torch
import torch.nn as nn
class PtnChainedModelWithAttentionMaskDef(nn.Module):
"""
A chained model that sequentially applies `model_a` and `model_b` optionally followed by a closure function. The model also takes attention mask for input.
### Fields:
- `model_a`: First model in the chain.
- `model_b`: Second model in the chain.
- `model_a_mask_req`: True if model_a requires attention_mask for input. For example, Set True if you plug in the `Ptn Multihead Attention` node.
- `model_b_mask_req`: True if model_b requires attention_mask for input
- `closure`: Optional differentiable function to be called at the end of forward.
pragma: skip_doc
"""
def __init__(self,
model_a: nn.Module,
model_b: nn.Module,
model_a_mask_req: bool,
model_b_mask_req: bool,
closure: Optional[Callable] = None) -> None:
super().__init__()
self.model_a = model_a
self.model_b = model_b
self.model_a_mask_req = model_a_mask_req
self.model_b_mask_req = model_b_mask_req
self.closure = closure
def forward(self, inputs: torch.Tensor, mask:torch.Tensor) -> torch.Tensor:
"""
Forward pass through the chained models.
Args:
inputs (torch.Tensor): Input tensor.
mask (torch.Tensor): Mask tensor.
Returns:
torch.Tensor: Output tensor after passing through `model_a`, `model_b`, and optionally `closure`.
"""
if self.model_a_mask_req:
x = self.model_a(inputs, mask)
else:
x = self.model_a(inputs)
if self.model_b_mask_req:
x = self.model_b(x, mask)
else:
x = self.model_b(x)
if self.closure:
x = self.closure(x)
return x
class PtnChainedModelWithAttentionMask:
"""
A chained model that sequentially applies `model_a` and `model_b`,
optionally followed by a closure function.
### Fields:
- `model_a`: First model in the chain.
- `model_b`: Second model in the chain.
- `model_a_mask_req`: True if model_a requires attention_mask for input
- `model_b_mask_req`: True if model_b requires attention_mask for input
- closure: Optional differentiable function to be called at the end of forward.
category: PyTorch wrapper - Model
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the expected input types.
Returns:
Dict[str, Any]: A dictionary specifying required and optional input types.
"""
return {
"required": {
"model_a": ("PTMODEL", {}),
"model_b": ("PTMODEL", {}),
"model_a_mask_req": ("BOOLEAN", {"default": False}),
"model_b_mask_req": ("BOOLEAN", {"default": False}),
},
"optional": {
"closure": ("PTCALLABLE", {})
}
}
RETURN_TYPES: Tuple[str, ...] = ("PTMODEL",)
FUNCTION: str = "f"
CATEGORY: str = "Training"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self,
model_a: nn.Module,
model_b: nn.Module,
model_a_mask_req: bool,
model_b_mask_req: bool,
closure: Optional[Callable] = None) -> Tuple[nn.Module]:
"""
Constructs a chained model with optional closure.
Args:
model_a (nn.Module): First model in the chain.
model_b (nn.Module): Second model in the chain.
model_a_mask_req (bool): Whether model_a needs attention mask.
model_b_mask_req (bool): Whether model_b needs attention mask.
closure (Callable, optional): A function applied after `model_b`.
Returns:
Tuple[nn.Module]: A tuple containing the instantiated model.
"""
with torch.inference_mode(False):
model = PtnChainedModelWithAttentionMaskDef(
model_a=model_a,
model_b=model_b,
model_a_mask_req=model_a_mask_req,
model_b_mask_req=model_b_mask_req,
closure=closure
)
return (model,)
modules/pytorch_wrapper/ptn_embedding.py
+1 -1
View File
@@ -26,7 +26,7 @@ class PtnEmbedding:
return {
"required": {
"num_embeddings": ("INT", {"default": 10000, "min": 1, "max": 1e6}),
"embedding_size": ("INT", {"default": 100, "min": 1, "max": 1e6}),
"embedding_dim": ("INT", {"default": 100, "min": 1, "max": 1e6}),
"padding_idx": ("INT", {"default": -1, "min": -1, "max": 1e6}),
}
}
modules/pytorch_wrapper/ptn_masked_mean_pooling.py
+77
View File
@@ -0,0 +1,77 @@
from typing import Any, Dict, Tuple
import torch
import torch.nn as nn
class PtnMaskedMeanPoolingDef(nn.Module):
"""
Constructs a masked mean pooling layer.
pragma: skip_doc
"""
def __init__(self) -> None:
super().__init__()
def forward(self, inputs: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
"""
Computes the mean of the input excluding the padding values using the given mask.
Args:
inputs (Tensor): 3D input tensor in (Batch, Sequence, Features) format.
mask (Tensor): 2D input tensor in (Batch, Sequence) format. Padding is indicated by 0; valid tokens are indicated by 1.
Returns:
torch.Tensor: Output tensor.
"""
if mask.dim() != 2:
raise ValueError("Mask is not rank 2")
if inputs.dim() != 3:
raise ValueError("Inputs is not rank 2")
maskf = mask.float() # (B, Seq)
x = (inputs * maskf.unsqueeze(-1)).sum(dim=1) / maskf.sum(dim=1, keepdim=True).clamp(min=1e-9) # (B, Feature)
if x.dim() != 2:
raise ValueError("Outputs is not rank 2")
return x
class PtnMaskedMeanPooling:
"""
Constructs a masked mean pooling layer.
Computes the mean of the input excluding the padding values using the given mask.
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 {
}
RETURN_TYPES: Tuple[str] = ("PTMODEL",)
FUNCTION: str = "f"
CATEGORY: str = "Training"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self) -> Tuple[nn.Module]:
"""
Constructs a masked mean pooling layer.
Computes the mean of the input excluding the padding values using the given mask.
Returns:
Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
"""
with torch.inference_mode(False):
model = PtnMaskedMeanPoolingDef()
return (model,)
modules/pytorch_wrapper/ptn_model_with_closure.py
+87
View File
@@ -0,0 +1,87 @@
from typing import Any, Dict, Tuple, Callable, Optional
import torch
import torch.nn as nn
class PtnModelWithClosureDef(nn.Module):
"""
A model that is followed by a closure function.
### Fields:
- `model`: The model in the chain.
- `closure`: Differentiable function to be called at the end of forward.
pragma: skip_doc
"""
def __init__(self,
model: nn.Module,
closure: Callable) -> None:
super().__init__()
self.model = model
self.closure = closure
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
"""
A model that is followed by a closure function.
Args:
- `model`: The model in the chain.
- `closure`: Differentiable function to be called at the end of forward.
Returns:
torch.Tensor: Output tensor after passing through `model` and `closure`.
"""
x = self.model(inputs)
x = self.closure(x)
return x
class PtnModelWithClosure:
"""
A model that is followed by a closure function.
### Fields:
- `model`: The model in the chain.
- `closure`: Differentiable function to be called at the end of forward. category: PyTorch wrapper - Model
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the expected input types.
Returns:
Dict[str, Any]: A dictionary specifying required and optional input types.
"""
return {
"required": {
"model": ("PTMODEL", {}),
"closure": ("PTCALLABLE", {})
}
}
RETURN_TYPES: Tuple[str, ...] = ("PTMODEL",)
FUNCTION: str = "f"
CATEGORY: str = "Training"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self,
model: nn.Module,
closure: Callable) -> Tuple[nn.Module]:
"""
A model that is followed by a closure function.
Args:
- `model`: The model in the chain.
- `closure`: Differentiable function to be called at the end of forward.
Returns:
Tuple[nn.Module]: A tuple containing the instantiated model.
"""
with torch.inference_mode(False):
model = PtnModelWithClosureDef(model, closure)
return (model,)
modules/pytorch_wrapper/ptn_multihead_attention.py
+153
View File
@@ -0,0 +1,153 @@
from typing import Any, Dict, Tuple
import torch
import torch.nn as nn
from typing import Tuple
class PtnMultiheadAttentionDef(nn.Module):
"""
Implements MultiheadAttention wrapper.
Attributes:
embed_dim (int): Hidden dimension (d_model) for each token of the Transformer.
num_heads (int): Number of attention heads.
dropout (float): Dropout rate for output.
bias (bool): If True, bias will be added to the linear projections of query, key, value, and the final output.
add_bias_kv (bool): If True, bias will be added to k and v sequences.
add_zero_attn (bool): If True, a zero batch will be added to k and v sequences.
kdim (int): Dimension of k.
vdim (int): Dimension of v.
batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature).
pragma: skip_doc
"""
def __init__(self,
embed_dim: int,
num_heads: int,
dropout: float,
bias: bool,
add_bias_kv: bool,
add_zero_attn: bool,
kdim: int,
vdim: int,
batch_first: bool
) -> None:
super().__init__()
self.model = nn.MultiheadAttention(
embed_dim=embed_dim,
num_heads=num_heads,
dropout=dropout,
bias=bias,
add_bias_kv=add_bias_kv,
add_zero_attn=add_zero_attn,
kdim=kdim,
vdim=vdim,
batch_first=batch_first
)
def forward(self, inputs: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
"""
Args:
inputs (torch.Tensor): Input tensor of shape (batch_size, seq_len, embed_dim) if batch_first is True.
mask (torch.Tensor): Attention mask where 1 indicates a valid token and 0 indicates padding.
This is the inverse of PyTorch's expected key_padding_mask, which masks out positions where the value is True.
Returns:
torch.Tensor: Output tensor after multi-head self-attention.
"""
return self.model(
inputs, inputs, inputs,
need_weights=False,
key_padding_mask=~mask.to(torch.bool)
)[0]
class PtnMultiheadAttention:
"""
A Multihead attention model.
Args:
embed_dim (int): Hidden dimension (d_model) for each token of the Transformer.
num_heads (int): Number of attention heads.
dropout (float): Dropout rate for output.
bias (bool): If True, bias will be added to the linear projections of query, key, value, and the final output.
add_bias_kv (bool): If True, bias will be added to k and v sequences.
add_zero_attn (bool): If True, a zero batch will be added to k and v sequences.
kdim (int): Dimension of k.
vdim (int): Dimension of v.
batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature).
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": {
"embed_dim": ("INT", {"default": 512, "min": 1, "max": 1e6}),
"num_heads": ("INT", {"default": 8, "min": 1, "max": 256}),
"dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1}),
"bias": ("BOOLEAN", {"default": True}),
"add_bias_kv": ("BOOLEAN", {"default": False}),
"add_zero_attn": ("BOOLEAN", {"default": False}),
"kdim": ("INT", {"default": 512, "min": 1, "max": 1e6}),
"vdim": ("INT", {"default": 512, "min": 1, "max": 1e6}),
"batch_first": ("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,
embed_dim: int,
num_heads: int,
dropout: float,
bias: bool,
add_bias_kv: bool,
add_zero_attn: bool,
kdim: int,
vdim: int,
batch_first: bool
) -> Tuple[nn.Module]:
"""
Constructs a Multihead attention model.
Args:
embed_dim (int): Hidden dimension (d_model) for each token of the Transformer.
num_heads (int): Number of attention heads.
dropout (float): Dropout rate for output.
bias (bool): If True, bias will be added to the linear projections of query, key, value, and the final output.
add_bias_kv (bool): If True, a learnable bias vector is prepended to the key and value sequences along the sequence dimension. Each bias vector has the same dimension as other sequence positions. Set this to False if you are unsure.
add_zero_attn (bool): If True, a zero vector will be added to k and v sequences.
kdim (int): Dimension of k.
vdim (int): Dimension of v.
batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature).
Returns:
Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
"""
with torch.inference_mode(False):
model = PtnMultiheadAttentionDef(
embed_dim=embed_dim,
num_heads=num_heads,
dropout=dropout,
bias=bias,
add_bias_kv=add_bias_kv,
add_zero_attn=add_zero_attn,
kdim=kdim,
vdim=vdim,
batch_first=batch_first
)
return (model,)
modules/pytorch_wrapper/ptn_residual_connection_model.py
+95
View File
@@ -0,0 +1,95 @@
from typing import Any, Dict, Tuple, Callable, Optional
import torch
import torch.nn as nn
class PtnResidualConnectionModelDef(nn.Module):
"""
A model that saves the input and add to the output of the specified model.
optionally followed by a closure function.
### Fields:
- `model`: A model to process input.
- `closure`: Optional differentiable function to be called at the end of forward.
pragma: skip_doc
"""
def __init__(self,
model: nn.Module,
closure: Optional[Callable] = None) -> None:
super().__init__()
self.model = model
self.closure = closure
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
"""
Forward pass through the chained models.
Args:
inputs (torch.Tensor): Input tensor.
Returns:
torch.Tensor: Output tensor after passing through `model` and adding residual input followed optionally by `closure`.
"""
res = inputs
x = self.model(inputs)
x = res + x
if self.closure:
x = self.closure(x)
return x
class PtnResidualConnectionModel:
"""
A model that saves the input and add to the output of the specified model.
optionally followed by a closure function.
### Fields:
- `model`: A model to process input.
- `closure`: Optional differentiable function to be called at the end of forward.
category: PyTorch wrapper - Model
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the expected input types.
Returns:
Dict[str, Any]: A dictionary specifying required and optional input types.
"""
return {
"required": {
"model": ("PTMODEL", {})
},
"optional": {
"closure": ("PTCALLABLE", {})
}
}
RETURN_TYPES: Tuple[str, ...] = ("PTMODEL",)
FUNCTION: str = "f"
CATEGORY: str = "Training"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self,
model: nn.Module,
closure: Optional[Callable] = None) -> Tuple[nn.Module]:
"""
Constructs a chained model with optional closure.
Args:
model (nn.Module): A model that saves the input and add to the output of the specified model.
closure ( Optional[Callable]): An optional differentiable function to be called at the end of forward.
Returns:
Tuple[nn.Module]: A tuple containing the instantiated model.
"""
with torch.inference_mode(False):
model = PtnResidualConnectionModelDef(model, closure)
return (model,)
modules/pytorch_wrapper/ptn_residual_connection_model_with_attention_mask.py
+96
View File
@@ -0,0 +1,96 @@
from typing import Any, Dict, Tuple, Callable, Optional
import torch
import torch.nn as nn
class PtnResidualConnectionModelWithAttentionMaskDef(nn.Module):
"""
A model that saves the input and add to the output of the specified model.
optionally followed by a closure function.
### Fields:
- `model`: A model to process input.
- `closure`: Optional differentiable function to be called at the end of forward.
pragma: skip_doc
"""
def __init__(self,
model: nn.Module,
closure: Optional[Callable] = None) -> None:
super().__init__()
self.model = model
self.closure = closure
def forward(self, inputs: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
"""
Forward pass through the chained models.
Args:
inputs (torch.Tensor): Input tensor.
mask (torch.Tensor): Attention mask.
Returns:
torch.Tensor: Output tensor after passing through `model` and adding residual input followed optionally by `closure`.
"""
res = inputs
x = self.model(inputs, mask)
x = res + x
if self.closure:
x = self.closure(x)
return x
class PtnResidualConnectionModelWithAttentionMask:
"""
A model that saves the input and add to the output of the specified model.
optionally followed by a closure function.
### Fields:
- `model`: A model to process input.
- `closure`: Optional differentiable function to be called at the end of forward.
category: PyTorch wrapper - Model
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the expected input types.
Returns:
Dict[str, Any]: A dictionary specifying required and optional input types.
"""
return {
"required": {
"model": ("PTMODEL", {})
},
"optional": {
"closure": ("PTCALLABLE", {})
}
}
RETURN_TYPES: Tuple[str, ...] = ("PTMODEL",)
FUNCTION: str = "f"
CATEGORY: str = "Training"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self,
model: nn.Module,
closure: Optional[Callable] = None) -> Tuple[nn.Module]:
"""
Constructs a chained model with optional closure.
Args:
model (nn.Module): A model that saves the input and add to the output of the specified model.
closure ( Optional[Callable]): An optional differentiable function to be called at the end of forward.
Returns:
Tuple[nn.Module]: A tuple containing the instantiated model.
"""
with torch.inference_mode(False):
model = PtnResidualConnectionModelWithAttentionMaskDef(model, closure)
return (model,)
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.6.1"
version = "4.0.0"
license = {file = "LICENSE"}
[project.urls]
tests/pytorch_wrapper/ptn_chained_model_with_attention_mask_test.py
+50
View File
@@ -0,0 +1,50 @@
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_chained_model_with_attention_mask import PtnChainedModelWithAttentionMask
from pytorch_wrapper.ptn_multihead_attention import PtnMultiheadAttention
class TestPtnChainedModel(unittest.TestCase):
def setUp(self) -> None:
"""Set up test instance with dummy models."""
self.node = PtnChainedModelWithAttentionMask()
self.model_a = PtnMultiheadAttention().f(
embed_dim=16,
num_heads=4,
dropout=0.1,
bias=True,
add_bias_kv=False,
add_zero_attn=False,
kdim=16,
vdim=16,
batch_first=True
)[0]
self.model_b = nn.LayerNorm([16]) # Normalize within each token
self.closure = nn.ReLU() # Apply ReLU as the closure
def test_forward_pass(self) -> None:
"""Test forward pass through the chained model."""
model, = self.node.f(
self.model_a,
self.model_b,
True, # a takes mask
False, # b does not take mask
self.closure)
x = torch.rand(2, 5, 16)
attention_mask = torch.ones(2, 5)
output = model(x, attention_mask)
self.assertEqual(output.shape, (2, 5, 16))
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/ptn_masked_mean_pooling_test.py
+82
View File
@@ -0,0 +1,82 @@
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_masked_mean_pooling import PtnMaskedMeanPooling
class TestPtnMaskedMeanPooling(unittest.TestCase):
def setUp(self):
"""Set up masked mean pooling model instance."""
self.model = PtnMaskedMeanPooling().f()[0]
self.model.eval()
# Each case: (inputs, mask, expected_output)
self.test_cases = [
(
torch.tensor([
[[1.0, 2.0], [3.0, 4.0], [0.0, 0.0]],
[[5.0, 5.0], [0.0, 0.0], [0.0, 0.0]]
]),
torch.tensor([
[1, 1, 0],
[1, 0, 0]
]),
torch.tensor([
[(1.0 + 3.0)/2, (2.0 + 4.0)/2],
[5.0, 5.0]
])
),
(
torch.tensor([
[[2.0, 4.0], [6.0, 8.0]]
]),
torch.tensor([
[0, 0]
]),
torch.tensor([
[0.0, 0.0]
])
),
(
torch.tensor([
[[1.0, 1.0], [2.0, 2.0], [3.0, 3.0]]
]),
torch.tensor([
[1, 1, 1]
]),
torch.tensor([
[(1+2+3)/3, (1+2+3)/3]
])
)
]
def test_masked_mean_pooling(self):
"""Test masked mean pooling against expected results."""
for i, (inputs, mask, expected) in enumerate(self.test_cases):
with self.subTest(i=i):
output = self.model(inputs, mask)
self.assertTrue(torch.allclose(output, expected, atol=1e-5), f"Mismatch in case {i}: got {output}, expected {expected}")
def test_invalid_mask_rank(self):
"""Test model raises ValueError for non-2D mask."""
inputs = torch.randn(2, 3, 4)
invalid_mask = torch.ones(2, 3, 1)
with self.assertRaises(ValueError):
self.model(inputs, invalid_mask)
def test_invalid_input_rank(self):
"""Test model raises ValueError for non-3D input."""
inputs = torch.randn(2, 3)
mask = torch.ones(2, 3)
with self.assertRaises(ValueError):
self.model(inputs, mask)
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/ptn_model_with_closure_test.py
+57
View File
@@ -0,0 +1,57 @@
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_model_with_closure import PtnModelWithClosure
class TestPtnModelWithClosure(unittest.TestCase):
def setUp(self) -> None:
"""Set up test instance with dummy models."""
self.node = PtnModelWithClosure()
class ModelA(nn.Module):
def __init__(self):
super().__init__()
self.layer = nn.Linear(3, 3, bias=False)
self.layer.weight = nn.Parameter(torch.eye(3), requires_grad=True) # Identity
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.layer(x)
self.model_a = ModelA()
self.closure = nn.ReLU() # Apply ReLU as the closure
def test_forward_pass(self) -> None:
"""Test forward pass through the chained model."""
model, = self.node.f(self.model_a, self.closure)
# Input tensor
input_tensor = torch.tensor([[1.0, 0.0, -1.0]], dtype=torch.float32)
# Expected transformation:
# ModelA (Identity): [1, 0, -1]
# Closure (ReLU): [1, 0, 0]
expected_output = torch.tensor([[1.0, 0.0, 0.0]], dtype=torch.float32)
# Get output from the chained model
output = model(input_tensor)
# Ensure output is a tensor
self.assertIsInstance(output, torch.Tensor, "Output is not a tensor")
# Ensure output shape matches expected
self.assertEqual(output.shape, expected_output.shape, "Output shape mismatch")
# Ensure output values match expected
self.assertTrue(torch.allclose(output, expected_output, atol=1e-6),
f"Expected {expected_output}, got {output}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/ptn_multihead_attention_test.py
+84
View File
@@ -0,0 +1,84 @@
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_multihead_attention import PtnMultiheadAttention
from pytorch_wrapper.ptn_multihead_attention import PtnMultiheadAttentionDef
class TestPtnMultiheadAttention(unittest.TestCase):
def setUp(self) -> None:
"""Set up test instance."""
self.node = PtnMultiheadAttention()
def test_forward_pass(self) -> None:
"""Test forward pass of MultiheadAttention with valid input."""
model, = self.node.f(
embed_dim=16,
num_heads=4,
dropout=0.1,
bias=True,
add_bias_kv=False,
add_zero_attn=False,
kdim=16,
vdim=16,
batch_first=True
)
self.assertIsInstance(model, PtnMultiheadAttentionDef)
# Create dummy input: batch_size=2, seq_len=5, embed_dim=16
x = torch.rand(2, 5, 16)
attention_mask = torch.ones(2, 5)
attn_output = model(x, attention_mask)
self.assertEqual(attn_output.shape, (2, 5, 16))
def test_forward_pass_with_zero_attn(self) -> None:
"""Test forward pass with add_zero_attn enabled."""
model, = self.node.f(
embed_dim=16,
num_heads=4,
dropout=0.0,
bias=True,
add_bias_kv=False,
add_zero_attn=True,
kdim=16,
vdim=16,
batch_first=True
)
x = torch.rand(2, 5, 16)
attention_mask = torch.ones(2, 5)
attn_output = model(x, attention_mask)
self.assertEqual(attn_output.shape, (2, 5, 16))
def test_forward_pass_with_bias_kv(self) -> None:
"""Test forward pass with bias_kv enabled."""
model, = self.node.f(
embed_dim=16,
num_heads=4,
dropout=0.0,
bias=True,
add_bias_kv=True,
add_zero_attn=False,
kdim=16,
vdim=16,
batch_first=True
)
x = torch.rand(2, 5, 16)
attention_mask = torch.ones(2, 5)
attn_output = model(x, attention_mask)
self.assertEqual(attn_output.shape, (2, 5, 16))
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/ptn_residual_connection_model_test.py
+68
View File
@@ -0,0 +1,68 @@
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_residual_connection_model import PtnResidualConnectionModel
class TestPtnResidualConnectionModel(unittest.TestCase):
def setUp(self) -> None:
"""Set up test instance with dummy model."""
self.node = PtnResidualConnectionModel()
class DummyModel(nn.Module):
def __init__(self):
super().__init__()
self.layer = nn.Linear(3, 3, bias=False)
self.layer.weight = nn.Parameter(torch.eye(3) * 0.2, requires_grad=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.layer(x)
self.model = DummyModel()
self.closure = nn.ReLU()
def test_forward_pass(self) -> None:
"""Test forward pass with residual addition and closure."""
model, = self.node.f(self.model, self.closure)
input_tensor = torch.tensor([[1.0, 0.0, -1.0]], dtype=torch.float32)
# model(input) = 0.2 * input = [0.2, 0.0, -0.2]
# residual + model(input) = [1.2, 0.0, -1.2]
# ReLU: [1.2, 0.0, 0.0]
expected_output = torch.tensor([[1.2, 0.0, 0.0]], dtype=torch.float32)
output = model(input_tensor)
self.assertIsInstance(output, torch.Tensor)
self.assertEqual(output.shape, expected_output.shape)
self.assertTrue(torch.allclose(output, expected_output, atol=1e-6),
f"Expected {expected_output}, got {output}")
def test_forward_pass_no_closure(self) -> None:
"""Test forward pass with residual addition without closure."""
model, = self.node.f(self.model, None)
input_tensor = torch.tensor([[1.0, 0.0, -1.0]], dtype=torch.float32)
# model(input) = 0.2 * input = [0.2, 0.0, -0.2]
# residual + model(input) = [1.2, 0.0, -1.2]
expected_output = torch.tensor([[1.2, 0.0, -1.2]], dtype=torch.float32)
output = model(input_tensor)
self.assertIsInstance(output, torch.Tensor)
self.assertEqual(output.shape, expected_output.shape)
self.assertTrue(torch.allclose(output, expected_output, atol=1e-6),
f"Expected {expected_output}, got {output}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/ptn_residual_connection_model_with_attention_mask_test.py
+41
View File
@@ -0,0 +1,41 @@
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_residual_connection_model_with_attention_mask import PtnResidualConnectionModelWithAttentionMask
from pytorch_wrapper.ptn_multihead_attention import PtnMultiheadAttention
class TestPtnResidualConnectionModelWithAttentionMask(unittest.TestCase):
def setUp(self) -> None:
"""Set up test instance with dummy model."""
self.node = PtnResidualConnectionModelWithAttentionMask()
self.inner_model = PtnMultiheadAttention().f(
embed_dim=16,
num_heads=4,
dropout=0.1,
bias=True,
add_bias_kv=False,
add_zero_attn=False,
kdim=16,
vdim=16,
batch_first=True
)[0]
self.closure = nn.GELU()
def test_forward_pass(self) -> None:
"""Test forward pass with residual addition and closure."""
model, = self.node.f(self.inner_model, self.closure)
x = torch.rand(2, 5, 16)
attention_mask = torch.ones(2, 5)
output = model(x, attention_mask)
self.assertEqual(output.shape, (2, 5, 16))
if __name__ == "__main__":
unittest.main()