diff --git a/__init__.py b/__init__.py index 0c06e6c..2e341d3 100644 --- a/__init__.py +++ b/__init__.py @@ -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", diff --git a/dev/scripts/reference_generate.py b/dev/scripts/reference_generate.py index a86053a..4a3c7df 100644 --- a/dev/scripts/reference_generate.py +++ b/dev/scripts/reference_generate.py @@ -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: diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index 537f308..e64618d 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -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 | diff --git a/docs/reference/pt_train_classification_transformer_model.md b/docs/reference/pt_train_classification_transformer_model.md index a7259e5..4514208 100644 --- a/docs/reference/pt_train_classification_transformer_model.md +++ b/docs/reference/pt_train_classification_transformer_model.md @@ -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 | diff --git a/docs/reference/ptn_chained_model_with_attention_mask.md b/docs/reference/ptn_chained_model_with_attention_mask.md new file mode 100644 index 0000000..8b66aa0 --- /dev/null +++ b/docs/reference/ptn_chained_model_with_attention_mask.md @@ -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 | + +
+Category: PyTorch wrapper - Model + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptn_embedding.md b/docs/reference/ptn_embedding.md index a14b2ac..60fbd18 100644 --- a/docs/reference/ptn_embedding.md +++ b/docs/reference/ptn_embedding.md @@ -10,7 +10,7 @@ Constructs an embedding layer. | Name | Data type | |---|---| | num_embeddings | Int | -| embedding_size | Int | +| embedding_dim | Int | | padding_idx | Int | ## Output diff --git a/docs/reference/ptn_masked_mean_pooling.md b/docs/reference/ptn_masked_mean_pooling.md new file mode 100644 index 0000000..6dee689 --- /dev/null +++ b/docs/reference/ptn_masked_mean_pooling.md @@ -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 | + +
+Category: PyTorch wrapper - Model + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptn_model_with_closure.md b/docs/reference/ptn_model_with_closure.md new file mode 100644 index 0000000..9169a6b --- /dev/null +++ b/docs/reference/ptn_model_with_closure.md @@ -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 | + +
+Category: Miscellaneous + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptn_multihead_attention.md b/docs/reference/ptn_multihead_attention.md new file mode 100644 index 0000000..92fa235 --- /dev/null +++ b/docs/reference/ptn_multihead_attention.md @@ -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 | + +
+Category: PyTorch wrapper - Model + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptn_residual_connection_model.md b/docs/reference/ptn_residual_connection_model.md new file mode 100644 index 0000000..cb7a6ce --- /dev/null +++ b/docs/reference/ptn_residual_connection_model.md @@ -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 | + +
+Category: PyTorch wrapper - Model + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptn_residual_connection_model_with_attention_mask.md b/docs/reference/ptn_residual_connection_model_with_attention_mask.md new file mode 100644 index 0000000..ee041cf --- /dev/null +++ b/docs/reference/ptn_residual_connection_model_with_attention_mask.md @@ -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 | + +
+Category: PyTorch wrapper - Model + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptv_hf_local_dataset.md b/docs/reference/ptv_hf_local_dataset.md index e5331c5..2f59b1f 100644 --- a/docs/reference/ptv_hf_local_dataset.md +++ b/docs/reference/ptv_hf_local_dataset.md @@ -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. diff --git a/examples/workflows/mha.json b/examples/workflows/mha.json new file mode 100644 index 0000000..3d4d4ca --- /dev/null +++ b/examples/workflows/mha.json @@ -0,0 +1 @@ +{"last_node_id":48,"last_link_id":85,"nodes":[{"id":24,"type":"SpLoadModel","pos":[49.8591194152832,3709.020263671875],"size":[210,60.936126708984375],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[{"name":"spmodel","localized_name":"spmodel","type":"SPMODEL","links":[44],"slot_index":0}],"properties":{"Node name for S&R":"SpLoadModel"},"widgets_values":["spiece.model"]},{"id":26,"type":"PtvHfDatasetWithTokenEncode","pos":[714.573974609375,3818.830810546875],"size":[271.5191650390625,178],"flags":{},"order":14,"mode":0,"inputs":[{"name":"encode","localized_name":"encode","type":"PTCALLABLE","link":38}],"outputs":[{"name":"PTVDATASET","localized_name":"PTVDATASET","type":"PTVDATASET","links":[40],"slot_index":0}],"properties":{"Node name for S&R":"PtvHfDatasetWithTokenEncode"},"widgets_values":["imdb","test","text","label",true,false]},{"id":27,"type":"PtnBCEWithLogitsLoss","pos":[1640.1048583984375,3897.827392578125],"size":[315,58],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[{"name":"PTLOSS","localized_name":"PTLOSS","type":"PTLOSS","links":[48],"slot_index":0}],"properties":{"Node name for S&R":"PtnBCEWithLogitsLoss"},"widgets_values":["mean"]},{"id":29,"type":"PtDataLoader","pos":[1093.905029296875,3812.4873046875],"size":[271.68658447265625,173.3039093017578],"flags":{},"order":18,"mode":0,"inputs":[{"name":"dataset","localized_name":"dataset","type":"PTVDATASET","link":40}],"outputs":[{"name":"PTDATALOADER","localized_name":"PTDATALOADER","type":"PTDATALOADER","links":[50],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[32,false,"{\"num_workers\":1}"]},{"id":35,"type":"PtTrainClassificationTransformerModel","pos":[2061.835693359375,3833.805908203125],"size":[344.3999938964844,302],"flags":{},"order":29,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":83},{"name":"train_loader","localized_name":"train_loader","type":"PTDATALOADER","link":46},{"name":"optimizer","localized_name":"optimizer","type":"PTOPTIMIZER","link":47},{"name":"loss_function","localized_name":"loss_function","type":"PTLOSS","link":48},{"name":"scheduler","localized_name":"scheduler","type":"PTLRSCHEDULER","shape":7,"link":49},{"name":"val_loader","localized_name":"val_loader","type":"PTDATALOADER","shape":7,"link":50}],"outputs":[{"name":"Model","localized_name":"Model","type":"PTMODEL","links":[41],"slot_index":0},{"name":"train loss","localized_name":"train loss","type":"TENSOR","links":null},{"name":"val loss","localized_name":"val loss","type":"TENSOR","links":null}],"properties":{"Node name for S&R":"PtTrainClassificationTransformerModel"},"widgets_values":[5,true,false,10,false,true,true]},{"id":20,"type":"PtnEmbedding","pos":[1910.9560546875,762.6825561523438],"size":[315,106],"flags":{},"order":12,"mode":0,"inputs":[{"name":"embedding_dim","type":"INT","widget":{"name":"embedding_dim"},"link":58}],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[60],"slot_index":0}],"properties":{"Node name for S&R":"PtnEmbedding"},"widgets_values":[32000,256,-1],"color":"#aeadaa","bgcolor":"#9a9996"},{"id":41,"type":"PtnChainedModelWithAttentionMask","pos":[2526.189208984375,932.9444580078125],"size":[310.79998779296875,144.12677001953125],"flags":{},"order":22,"mode":0,"inputs":[{"name":"model_a","localized_name":"model_a","type":"PTMODEL","link":60},{"name":"model_b","localized_name":"model_b","type":"PTMODEL","link":61},{"name":"closure","localized_name":"closure","type":"PTCALLABLE","shape":7,"link":null}],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[62],"slot_index":0}],"properties":{"Node name for S&R":"PtnChainedModelWithAttentionMask"},"widgets_values":[false,true],"color":"#432","bgcolor":"#653"},{"id":44,"type":"PtnResidualConnectionModel","pos":[1731.1656494140625,1890.8028564453125],"size":[243.60000610351562,54.06667709350586],"flags":{},"order":21,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":74},{"name":"closure","localized_name":"closure","type":"PTCALLABLE","shape":7,"link":null}],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[77],"slot_index":0}],"properties":{"Node name for S&R":"PtnResidualConnectionModel"},"widgets_values":[],"color":"#222","bgcolor":"#000"},{"id":8,"type":"PtnChainedModel","pos":[1494.827880859375,1888.4935302734375],"size":[174.87501525878906,67.0083236694336],"flags":{},"order":19,"mode":0,"inputs":[{"name":"model_a","localized_name":"model_a","type":"PTMODEL","link":18},{"name":"model_b","localized_name":"model_b","type":"PTMODEL","link":7},{"name":"closure","localized_name":"closure","type":"PTCALLABLE","shape":7,"link":null}],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[74],"slot_index":0}],"properties":{"Node name for S&R":"PtnChainedModel"},"widgets_values":[],"color":"#432","bgcolor":"#653"},{"id":16,"type":"PtnModelWithClosure","pos":[1206.2276611328125,1807.6285400390625],"size":[200.566650390625,56.18415069580078],"flags":{},"order":15,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":16},{"name":"closure","localized_name":"closure","type":"PTCALLABLE","link":17}],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[18],"slot_index":0}],"properties":{"Node name for S&R":"PtnModelWithClosure"},"widgets_values":[],"color":"#332922","bgcolor":"#593930"},{"id":3,"type":"PtnChainedModelWithAttentionMask","pos":[1867.1180419921875,1286.6824951171875],"size":[333.6241455078125,122],"flags":{},"order":20,"mode":0,"inputs":[{"name":"model_a","localized_name":"model_a","type":"PTMODEL","link":28},{"name":"model_b","localized_name":"model_b","type":"PTMODEL","link":3},{"name":"closure","localized_name":"closure","type":"PTCALLABLE","shape":7,"link":null}],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[61],"slot_index":0}],"properties":{"Node name for S&R":"PtnChainedModelWithAttentionMask"},"widgets_values":[true,false],"color":"#432","bgcolor":"#653"},{"id":1,"type":"PtnResidualConnectionModelWithAttentionMask","pos":[1405.3084716796875,1099.3072509765625],"size":[411.6000061035156,91.88177490234375],"flags":{},"order":16,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":27},{"name":"closure","localized_name":"closure","type":"PTCALLABLE","shape":7,"link":null}],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[28],"slot_index":0}],"properties":{"Node name for S&R":"PtnResidualConnectionModelWithAttentionMask"},"widgets_values":[],"color":"#222","bgcolor":"#000"},{"id":4,"type":"PtnLayerNorm","pos":[1443.7354736328125,1456.87890625],"size":[307.50115966796875,107.51513671875],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[3],"slot_index":0}],"properties":{"Node name for S&R":"PtnLayerNorm"},"widgets_values":["[256]",true,true],"color":"#232","bgcolor":"#353"},{"id":10,"type":"PtnLayerNorm","pos":[1734.420166015625,2026.943359375],"size":[251.03451538085938,106],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[71],"slot_index":0}],"properties":{"Node name for S&R":"PtnLayerNorm"},"widgets_values":["[256]",true,true],"color":"#232","bgcolor":"#353"},{"id":6,"type":"PtnLinear","pos":[948.1487426757812,1770.8453369140625],"size":[210,108.21834564208984],"flags":{},"order":8,"mode":0,"inputs":[{"name":"in_features","type":"INT","widget":{"name":"in_features"},"link":19}],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[16],"slot_index":0}],"properties":{"Node name for S&R":"PtnLinear"},"widgets_values":[512,1024,true],"color":"#223","bgcolor":"#335"},{"id":7,"type":"PtnLinear","pos":[1215.0638427734375,2012.450927734375],"size":[210,110.95808410644531],"flags":{},"order":9,"mode":0,"inputs":[{"name":"out_features","type":"INT","widget":{"name":"out_features"},"link":20}],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[7],"slot_index":0}],"properties":{"Node name for S&R":"PtnLinear"},"widgets_values":[1024,512,true],"color":"#223","bgcolor":"#335"},{"id":2,"type":"PtnMultiheadAttention","pos":[1018.49951171875,1097.602783203125],"size":[315,290],"flags":{},"order":10,"mode":0,"inputs":[{"name":"embed_dim","type":"INT","widget":{"name":"embed_dim"},"link":21},{"name":"kdim","type":"INT","widget":{"name":"kdim"},"link":22},{"name":"vdim","type":"INT","widget":{"name":"vdim"},"link":23}],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[27],"slot_index":0}],"properties":{"Node name for S&R":"PtnMultiheadAttention"},"widgets_values":[512,8,0,true,false,false,512,512,true],"color":"#322","bgcolor":"#533"},{"id":19,"type":"CDAIntCreate","pos":[295.0705261230469,1197.1085205078125],"size":[210,72.17996978759766],"flags":{},"order":4,"mode":0,"inputs":[],"outputs":[{"name":"INT","localized_name":"INT","type":"INT","links":[19,20,21,22,23,55,58],"slot_index":0}],"title":"Embedding dimension","properties":{"Node name for S&R":"CDAIntCreate"},"widgets_values":["256"]},{"id":22,"type":"PtnChainedModelWithAttentionMask","pos":[2526.189208984375,1597.4144287109375],"size":[310.79998779296875,144.12677001953125],"flags":{},"order":24,"mode":0,"inputs":[{"name":"model_a","localized_name":"model_a","type":"PTMODEL","link":62},{"name":"model_b","localized_name":"model_b","type":"PTMODEL","link":78},{"name":"closure","localized_name":"closure","type":"PTCALLABLE","shape":7,"link":null}],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[63],"slot_index":0}],"properties":{"Node name for S&R":"PtnChainedModelWithAttentionMask"},"widgets_values":[true,false],"color":"#432","bgcolor":"#653"},{"id":43,"type":"PtnChainedModel","pos":[2057.467041015625,1958.353515625],"size":[178.70663452148438,68.2183837890625],"flags":{},"order":23,"mode":0,"inputs":[{"name":"model_a","localized_name":"model_a","type":"PTMODEL","link":77},{"name":"model_b","localized_name":"model_b","type":"PTMODEL","link":71},{"name":"closure","localized_name":"closure","type":"PTCALLABLE","shape":7,"link":null}],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[78],"slot_index":0}],"properties":{"Node name for S&R":"PtnChainedModel"},"widgets_values":[],"color":"#432","bgcolor":"#653"},{"id":17,"type":"PtfGELU","pos":[1001.5574340820312,1969.317626953125],"size":[148.99588012695312,41.5283088684082],"flags":{},"order":5,"mode":0,"inputs":[],"outputs":[{"name":"PTCALLABLE","localized_name":"PTCALLABLE","type":"PTCALLABLE","links":[17],"slot_index":0}],"properties":{"Node name for S&R":"PtfGELU"},"widgets_values":[],"color":"#332922","bgcolor":"#593930"},{"id":36,"type":"PtnLinear","pos":[1888.2996826171875,2754.08447265625],"size":[315,106],"flags":{},"order":11,"mode":0,"inputs":[{"name":"in_features","type":"INT","widget":{"name":"in_features"},"link":55}],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[84],"slot_index":0}],"properties":{"Node name for S&R":"PtnLinear"},"widgets_values":[1,1,true],"color":"#223","bgcolor":"#335"},{"id":47,"type":"PtnMaskedMeanPooling","pos":[1930.5147705078125,2421.533935546875],"size":[289.79998779296875,26],"flags":{},"order":6,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[85],"slot_index":0}],"properties":{"Node name for S&R":"PtnMaskedMeanPooling"},"widgets_values":[]},{"id":28,"type":"PtDataLoader","pos":[1096.95263671875,3531.836181640625],"size":[271.68658447265625,173.3039093017578],"flags":{},"order":17,"mode":0,"inputs":[{"name":"dataset","localized_name":"dataset","type":"PTVDATASET","link":39}],"outputs":[{"name":"PTDATALOADER","localized_name":"PTDATALOADER","type":"PTDATALOADER","links":[46],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[32,true,"{\"num_workers\":1}"]},{"id":32,"type":"PtoAdamW","pos":[1539.4169921875,3517.294189453125],"size":[285.66900634765625,158.52035522460938],"flags":{},"order":27,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":82}],"outputs":[{"name":"PTOPTIMIZER","localized_name":"PTOPTIMIZER","type":"PTOPTIMIZER","links":[42,47],"slot_index":0}],"properties":{"Node name for S&R":"PtoAdamW"},"widgets_values":[0.0002999806000000002,0.9,0.9989999956000005,0.01,false]},{"id":31,"type":"PtoLrSchedulerCosineAnnealing","pos":[1991.636474609375,3522.379638671875],"size":[311.11163330078125,113.72386932373047],"flags":{},"order":28,"mode":0,"inputs":[{"name":"optimizer","localized_name":"optimizer","type":"PTOPTIMIZER","link":42}],"outputs":[{"name":"PTLRSCHEDULER","localized_name":"PTLRSCHEDULER","type":"PTLRSCHEDULER","links":[49],"slot_index":0}],"properties":{"Node name for S&R":"PtoLrSchedulerCosineAnnealing"},"widgets_values":[10,1.0000000000000005e-9]},{"id":39,"type":"PtnChainedModelWithAttentionMask","pos":[2526.189208984375,2213.20068359375],"size":[310.79998779296875,127.73847198486328],"flags":{},"order":25,"mode":0,"inputs":[{"name":"model_a","localized_name":"model_a","type":"PTMODEL","link":63},{"name":"model_b","localized_name":"model_b","type":"PTMODEL","link":85},{"name":"closure","localized_name":"closure","type":"PTCALLABLE","shape":7,"link":null}],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[79],"slot_index":0}],"properties":{"Node name for S&R":"PtnChainedModelWithAttentionMask"},"widgets_values":[true,true],"color":"#432","bgcolor":"#653"},{"id":34,"type":"SpEncode","pos":[306.2023010253906,3710.096923828125],"size":[315,154],"flags":{},"order":7,"mode":0,"inputs":[{"name":"spmodel","localized_name":"spmodel","type":"SPMODEL","link":44}],"outputs":[{"name":"PTCALLABLE","localized_name":"PTCALLABLE","type":"PTCALLABLE","links":[37,38],"slot_index":0}],"properties":{"Node name for S&R":"SpEncode"},"widgets_values":[true,"max_length",0,true,473]},{"id":23,"type":"PandasShowText","pos":[2629.6103515625,4024.010498046875],"size":[302.7994079589844,112],"flags":{},"order":31,"mode":0,"inputs":[{"name":"text","type":"STRING","widget":{"name":"text"},"link":36}],"outputs":[{"name":"STRING","localized_name":"STRING","type":"STRING","shape":6,"links":null}],"properties":{"Node name for S&R":"PandasShowText"},"widgets_values":["","/container/transfer/models/custom_transformer_model.pt"]},{"id":45,"type":"PtnChainedModelWithAttentionMask","pos":[2526.189208984375,2555.535888671875],"size":[310.79998779296875,129.08058166503906],"flags":{},"order":26,"mode":0,"inputs":[{"name":"model_a","localized_name":"model_a","type":"PTMODEL","link":79},{"name":"model_b","localized_name":"model_b","type":"PTMODEL","link":84},{"name":"closure","localized_name":"closure","type":"PTCALLABLE","shape":7,"link":null}],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[82,83],"slot_index":0}],"properties":{"Node name for S&R":"PtnChainedModelWithAttentionMask"},"widgets_values":[true,false],"color":"#432","bgcolor":"#653"},{"id":25,"type":"PtvHfDatasetWithTokenEncode","pos":[710.87939453125,3542.5537109375],"size":[271.5191650390625,178],"flags":{},"order":13,"mode":0,"inputs":[{"name":"encode","localized_name":"encode","type":"PTCALLABLE","link":37}],"outputs":[{"name":"PTVDATASET","localized_name":"PTVDATASET","type":"PTVDATASET","links":[39],"slot_index":0}],"properties":{"Node name for S&R":"PtvHfDatasetWithTokenEncode"},"widgets_values":["imdb","train","text","label",true,false]},{"id":30,"type":"PtSaveModel","pos":[2617.819580078125,3836.388671875],"size":[318.3275146484375,91.27499389648438],"flags":{},"order":30,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":41}],"outputs":[{"name":"model full path","localized_name":"model full path","type":"STRING","links":[36],"slot_index":0}],"properties":{"Node name for S&R":"PtSaveModel"},"widgets_values":["/container/transfer/models/custom_transformer_model.pt"]}],"links":[[3,4,0,3,1,"PTMODEL"],[7,7,0,8,1,"PTMODEL"],[16,6,0,16,0,"PTMODEL"],[17,17,0,16,1,"PTCALLABLE"],[18,16,0,8,0,"PTMODEL"],[19,19,0,6,0,"INT"],[20,19,0,7,0,"INT"],[21,19,0,2,0,"INT"],[22,19,0,2,1,"INT"],[23,19,0,2,2,"INT"],[27,2,0,1,0,"PTMODEL"],[28,1,0,3,0,"PTMODEL"],[36,30,0,23,0,"STRING"],[37,34,0,25,0,"PTCALLABLE"],[38,34,0,26,0,"PTCALLABLE"],[39,25,0,28,0,"PTVDATASET"],[40,26,0,29,0,"PTVDATASET"],[41,35,0,30,0,"PTMODEL"],[42,32,0,31,0,"PTOPTIMIZER"],[44,24,0,34,0,"SPMODEL"],[46,28,0,35,1,"PTDATALOADER"],[47,32,0,35,2,"PTOPTIMIZER"],[48,27,0,35,3,"PTLOSS"],[49,31,0,35,4,"PTLRSCHEDULER"],[50,29,0,35,5,"PTDATALOADER"],[55,19,0,36,0,"INT"],[58,19,0,20,0,"INT"],[60,20,0,41,0,"PTMODEL"],[61,3,0,41,1,"PTMODEL"],[62,41,0,22,0,"PTMODEL"],[63,22,0,39,0,"PTMODEL"],[71,10,0,43,1,"PTMODEL"],[74,8,0,44,0,"PTMODEL"],[77,44,0,43,0,"PTMODEL"],[78,43,0,22,1,"PTMODEL"],[79,39,0,45,0,"PTMODEL"],[82,45,0,32,0,"PTMODEL"],[83,45,0,35,0,"PTMODEL"],[84,36,0,45,1,"PTMODEL"],[85,47,0,39,1,"PTMODEL"]],"groups":[{"id":6,"title":"Transformer model","bounding":[58.61259078979492,496.19342041015625,2965.295166015625,2505.033203125],"color":"#444","font_size":24,"flags":{}},{"id":1,"title":"Constant definition","bounding":[191.04232788085938,1002.2987060546875,422.5166015625,397.8311767578125],"color":"#3f789e","font_size":24,"flags":{}},{"id":2,"title":"Token to Embedding Conversion","bounding":[1889.7581787109375,668.4641723632812,372.5622253417969,234.32557678222656],"color":"#c64600","font_size":24,"flags":{}},{"id":3,"title":"Attention Block","bounding":[916.3255004882812,1000.0462646484375,1354.99560546875,614.4207763671875],"color":"#26a269","font_size":24,"flags":{}},{"id":4,"title":"Feed forward block","bounding":[910.754638671875,1675.8887939453125,1355.9459228515625,502.7033386230469],"color":"#1a5fb4","font_size":24,"flags":{}},{"id":5,"title":"Linear head","bounding":[1859.5028076171875,2653.4169921875,395.3128967285156,262.0767822265625],"color":"#613583","font_size":24,"flags":{}},{"id":7,"title":"Masked Mean","bounding":[1863.8021240234375,2302.5185546875,395.3128967285156,262.0767822265625],"color":"#613583","font_size":24,"flags":{}},{"id":8,"title":"Data loading","bounding":[32.85782241821289,3425.364501953125,1389.021240234375,763.4718627929688],"color":"#26a269","font_size":24,"flags":{}},{"id":9,"title":"Trainer","bounding":[1522.4901123046875,3427.466064453125,913.7997436523438,758.6957397460938],"color":"#e01b24","font_size":24,"flags":{}},{"id":10,"title":"Model saving","bounding":[2571.20703125,3431.988037109375,438.21942138671875,743.6028442382812],"color":"#1a5fb4","font_size":24,"flags":{}}],"config":{},"extra":{"ds":{"scale":0.2091956454947031,"offset":[2201.2200384524076,-243.6224455379711]},"node_versions":{"pt-wrapper":"f19a6ebcfe21dad3834b9fd31636fd93a57eb608","data-analysis":"68dd090367adffa569c3b3993d56e422f4516d2d"},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/modules/pytorch_wrapper/pt_train_classification_transformer_model.py b/modules/pytorch_wrapper/pt_train_classification_transformer_model.py index f867ad5..034fe3f 100644 --- a/modules/pytorch_wrapper/pt_train_classification_transformer_model.py +++ b/modules/pytorch_wrapper/pt_train_classification_transformer_model.py @@ -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) diff --git a/modules/pytorch_wrapper/ptn_chained_model.py b/modules/pytorch_wrapper/ptn_chained_model.py index 8e27175..a6f0769 100644 --- a/modules/pytorch_wrapper/ptn_chained_model.py +++ b/modules/pytorch_wrapper/ptn_chained_model.py @@ -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 diff --git a/modules/pytorch_wrapper/ptn_chained_model_with_attention_mask.py b/modules/pytorch_wrapper/ptn_chained_model_with_attention_mask.py new file mode 100644 index 0000000..1a7f1fd --- /dev/null +++ b/modules/pytorch_wrapper/ptn_chained_model_with_attention_mask.py @@ -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,) \ No newline at end of file diff --git a/modules/pytorch_wrapper/ptn_embedding.py b/modules/pytorch_wrapper/ptn_embedding.py index a6be914..a1894de 100644 --- a/modules/pytorch_wrapper/ptn_embedding.py +++ b/modules/pytorch_wrapper/ptn_embedding.py @@ -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}), } } diff --git a/modules/pytorch_wrapper/ptn_masked_mean_pooling.py b/modules/pytorch_wrapper/ptn_masked_mean_pooling.py new file mode 100644 index 0000000..a86390c --- /dev/null +++ b/modules/pytorch_wrapper/ptn_masked_mean_pooling.py @@ -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,) diff --git a/modules/pytorch_wrapper/ptn_model_with_closure.py b/modules/pytorch_wrapper/ptn_model_with_closure.py new file mode 100644 index 0000000..ebe693b --- /dev/null +++ b/modules/pytorch_wrapper/ptn_model_with_closure.py @@ -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,) diff --git a/modules/pytorch_wrapper/ptn_multihead_attention.py b/modules/pytorch_wrapper/ptn_multihead_attention.py new file mode 100644 index 0000000..ca1d74f --- /dev/null +++ b/modules/pytorch_wrapper/ptn_multihead_attention.py @@ -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,) diff --git a/modules/pytorch_wrapper/ptn_residual_connection_model.py b/modules/pytorch_wrapper/ptn_residual_connection_model.py new file mode 100644 index 0000000..9ef9fd1 --- /dev/null +++ b/modules/pytorch_wrapper/ptn_residual_connection_model.py @@ -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,) diff --git a/modules/pytorch_wrapper/ptn_residual_connection_model_with_attention_mask.py b/modules/pytorch_wrapper/ptn_residual_connection_model_with_attention_mask.py new file mode 100644 index 0000000..ad2715d --- /dev/null +++ b/modules/pytorch_wrapper/ptn_residual_connection_model_with_attention_mask.py @@ -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,) diff --git a/pyproject.toml b/pyproject.toml index a1736ff..0eada01 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "pt-wrapper" description = "PyTorch extension for ComfyUI featuring extensive PyTorch wrapper nodes for seamless tensor operations and PyTorch model training." -version = "3.6.1" +version = "4.0.0" license = {file = "LICENSE"} [project.urls] diff --git a/tests/pytorch_wrapper/ptn_chained_model_with_attention_mask_test.py b/tests/pytorch_wrapper/ptn_chained_model_with_attention_mask_test.py new file mode 100644 index 0000000..4743707 --- /dev/null +++ b/tests/pytorch_wrapper/ptn_chained_model_with_attention_mask_test.py @@ -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() diff --git a/tests/pytorch_wrapper/ptn_masked_mean_pooling_test.py b/tests/pytorch_wrapper/ptn_masked_mean_pooling_test.py new file mode 100644 index 0000000..da56a0b --- /dev/null +++ b/tests/pytorch_wrapper/ptn_masked_mean_pooling_test.py @@ -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() diff --git a/tests/pytorch_wrapper/ptn_model_with_closure_test.py b/tests/pytorch_wrapper/ptn_model_with_closure_test.py new file mode 100644 index 0000000..db3a956 --- /dev/null +++ b/tests/pytorch_wrapper/ptn_model_with_closure_test.py @@ -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() diff --git a/tests/pytorch_wrapper/ptn_multihead_attention_test.py b/tests/pytorch_wrapper/ptn_multihead_attention_test.py new file mode 100644 index 0000000..3da9ba1 --- /dev/null +++ b/tests/pytorch_wrapper/ptn_multihead_attention_test.py @@ -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() diff --git a/tests/pytorch_wrapper/ptn_residual_connection_model_test.py b/tests/pytorch_wrapper/ptn_residual_connection_model_test.py new file mode 100644 index 0000000..3950f73 --- /dev/null +++ b/tests/pytorch_wrapper/ptn_residual_connection_model_test.py @@ -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() diff --git a/tests/pytorch_wrapper/ptn_residual_connection_model_with_attention_mask_test.py b/tests/pytorch_wrapper/ptn_residual_connection_model_with_attention_mask_test.py new file mode 100644 index 0000000..7206b9e --- /dev/null +++ b/tests/pytorch_wrapper/ptn_residual_connection_model_with_attention_mask_test.py @@ -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()