diff --git a/__init__.py b/__init__.py index 2e341d3..30b51a6 100644 --- a/__init__.py +++ b/__init__.py @@ -182,6 +182,7 @@ 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_multihead_attention_custom import PtnMultiheadAttentionCustom 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 @@ -402,6 +403,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtnMaxPool2d": PtnMaxPool2d, "PtnModelWithClosure": PtnModelWithClosure, "PtnMultiheadAttention": PtnMultiheadAttention, + "PtnMultiheadAttentionCustom": PtnMultiheadAttentionCustom, "PtnNLLLoss": PtnNLLLoss, "PtnPreAddChannelAxis": PtnPreAddChannelAxis, "PtnPreFlatten": PtnPreFlatten, @@ -622,6 +624,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtnMaxPool2d": "Ptn Max Pool 2d", "PtnModelWithClosure": "Ptn Model With Closure", "PtnMultiheadAttention": "Ptn Multihead Attention", + "PtnMultiheadAttentionCustom": "Ptn Multihead Attention Custom", "PtnNLLLoss": "Ptn NLL Loss", "PtnPreAddChannelAxis": "Ptn Pre Add Channel Axis", "PtnPreFlatten": "Ptn Pre Flatten", diff --git a/docs/building_transformer_from_scratch.md b/docs/building_transformer_from_scratch.md index a7d74a1..36245cd 100644 --- a/docs/building_transformer_from_scratch.md +++ b/docs/building_transformer_from_scratch.md @@ -167,3 +167,11 @@ With this workflow, you should reach around **85% validation accuracy** on the H This is a solid performance for a Transformer built entirely from scratch. --- + +### If you want to understand how Multi-head Attention is implemented + +Refer to `modules/pytorch_wrapper/ptn_multihead_attention_custom.py`, which contains the implementation of Multi-head Attention and its wrapper node. +You can also check out the example workflow in `examples/workflows/mha_custom.json`, which demonstrates how to use this wrapper node. +This workflow achieves around 85–86% val accuracy on the IMDB dataset. + +--- diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index e64618d..dec7510 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -176,6 +176,7 @@ | [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 Multihead Attention Custom](ptn_multihead_attention_custom.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. | diff --git a/docs/reference/ptn_multihead_attention_custom.md b/docs/reference/ptn_multihead_attention_custom.md new file mode 100644 index 0000000..539e9f0 --- /dev/null +++ b/docs/reference/ptn_multihead_attention_custom.md @@ -0,0 +1,32 @@ +# Ptn Multihead Attention Custom +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. + 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 | +| 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/examples/workflows/mha_custom.json b/examples/workflows/mha_custom.json new file mode 100644 index 0000000..0be8e66 --- /dev/null +++ b/examples/workflows/mha_custom.json @@ -0,0 +1 @@ +{"last_node_id":49,"last_link_id":89,"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":11,"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":89},{"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":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":4,"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":10,"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":5,"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"]},{"id":19,"type":"CDAIntCreate","pos":[295.0705261230469,1197.1085205078125],"size":[210,72.17996978759766],"flags":{},"order":6,"mode":0,"inputs":[],"outputs":[{"name":"INT","localized_name":"INT","type":"INT","links":[19,20,55,58,86,87,88],"slot_index":0}],"title":"Embedding dimension","properties":{"Node name for S&R":"CDAIntCreate"},"widgets_values":["256"]},{"id":49,"type":"PtnMultiheadAttentionCustom","pos":[997.123779296875,1152.76220703125],"size":[308.3500061035156,247.6350555419922],"flags":{},"order":12,"mode":0,"inputs":[{"name":"embed_dim","type":"INT","widget":{"name":"embed_dim"},"link":86},{"name":"kdim","type":"INT","widget":{"name":"kdim"},"link":87},{"name":"vdim","type":"INT","widget":{"name":"vdim"},"link":88}],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[89],"slot_index":0}],"properties":{"Node name for S&R":"PtnMultiheadAttentionCustom"},"widgets_values":[512,8,0,true,512,512,true]}],"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"],[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"],[86,19,0,49,0,"INT"],[87,19,0,49,1,"INT"],[88,19,0,49,2,"INT"],[89,49,0,1,0,"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.8133685800604232,"offset":[-158.43105621502323,-950.0036651358898]},"node_versions":{"pt-wrapper":"55ebe2293c7a09d610c27e8f95b46ac8317f4050","data-analysis":"68dd090367adffa569c3b3993d56e422f4516d2d"},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/modules/pytorch_wrapper/ptn_multihead_attention_custom.py b/modules/pytorch_wrapper/ptn_multihead_attention_custom.py new file mode 100644 index 0000000..7b4b86c --- /dev/null +++ b/modules/pytorch_wrapper/ptn_multihead_attention_custom.py @@ -0,0 +1,209 @@ +from typing import Any, Dict, Tuple, Optional +import math +import torch +import torch.nn as nn +from einops import einops + +class PtnMultiheadAttentionCustomDef(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, + kdim: int, + vdim: int, + batch_first: bool + ) -> None: + super().__init__() + + # Compute dimensions + self.embed_dim = embed_dim + self.num_heads = num_heads + self.head_dim = embed_dim // num_heads + assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads" + + # Projection matrix + self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + self.k_proj = nn.Linear(kdim, embed_dim, bias=bias) + self.v_proj = nn.Linear(vdim, embed_dim, bias=bias) + self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias) + + if dropout: + self.dropout = nn.Dropout(dropout) + else: + self.dropout = None + + self.batch_first = batch_first + + + def forward(self, + inputs: torch.Tensor, + mask:Optional[torch.Tensor]=None) -> torch.Tensor: + """ + Args: + inputs (torch.Tensor): Input tensor of shape (batch_size, seq_len, embed_dim) if batch_first is True. + mask:Optional[torch.Tensor]: An optional 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. + """ + if inputs.dim() != 3: + raise ValueError("Inputs is not a rank 3 tensor.") + + x = inputs + + if not self.batch_first: # (seq, batch, features) to (batch, seq, feature) + x = einops.rearrange(x, 's b f -> b s f') + + q = self.q_proj(x) + k = self.k_proj(x) + v = self.v_proj(x) + + # Split feature axis into heads × (feature per head). + # The size of this per-head feature is denoted by "d" (aka "head_dim"). + # Rearrange so that dimensions are (batch, heads, seq, d). + # (h d) means h * d. + q = einops.rearrange(q, 'b s (h d) -> b h s d', h=self.num_heads) + k = einops.rearrange(k, 'b s (h d) -> b h s d', h=self.num_heads) + v = einops.rearrange(v, 'b s (h d) -> b h s d', h=self.num_heads) + + # dot product q @ k.T / sqrt(head_dim) + kt = torch.transpose(k, -1, -2) + dot = torch.matmul(q, kt) / math.sqrt(self.head_dim) # torch.sqrt requires a tensor + + # q @ k = (b h s d) @ (b h d s) = (b h s s) + if dot.size() != (q.size(0), q.size(1), q.size(2), q.size(2)): + raise ValueError(f"Unexpected size for q@k.T: {dot.size()}") + + # Apply mask (1 = valid token, 0 = padding) + # If token contains 0, the attention score for the cell + # will be replaced with -inf, which will be converted to 0 + # in softmax. + # e.g. one row of attention score + # pos 0 1 2 ... 1023 + # 5 2 -5 -20 + # mask + # 1 1 0 0 + # result + # 5 2 -inf -inf + if mask is not None: + mask = mask.to(torch.bool) + mask = mask.unsqueeze(1).unsqueeze(2) # (B, 1, 1, S) + dot = dot.masked_fill(mask == 0, float('-inf')) + + # Normalize + dot = torch.softmax(dot, -1) + + # Apply attention weights to value vectors + v = torch.matmul(dot, v) + + # Rearrange + v = einops.rearrange(v, 'b h s d -> b s (h d)') + + x = self.out_proj(v) + + if self.dropout: + x = self.dropout(x) + + if not self.batch_first: + x = einops.rearrange(x, 'b s f -> s b f') + + return x + + +class PtnMultiheadAttentionCustom: + """ + 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. + 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}), + "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, + 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. + 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 = PtnMultiheadAttentionCustomDef( + embed_dim=embed_dim, + num_heads=num_heads, + dropout=dropout, + bias=bias, + kdim=kdim, + vdim=vdim, + batch_first=batch_first + ) + return (model,) diff --git a/pyproject.toml b/pyproject.toml index 3c1e6da..45f62b7 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 = "4.0.1" +version = "4.1.0" license = {file = "LICENSE"} [project.urls] diff --git a/tests/pytorch_wrapper/ptn_multihead_attention_custom_test.py b/tests/pytorch_wrapper/ptn_multihead_attention_custom_test.py new file mode 100644 index 0000000..e9c8a8b --- /dev/null +++ b/tests/pytorch_wrapper/ptn_multihead_attention_custom_test.py @@ -0,0 +1,62 @@ +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_custom import PtnMultiheadAttentionCustom +from pytorch_wrapper.ptn_multihead_attention_custom import PtnMultiheadAttentionCustomDef + +class TestPtnMultiheadAttentionCustom(unittest.TestCase): + def setUp(self) -> None: + """Set up test instance.""" + self.node = PtnMultiheadAttentionCustom() + + def test_fbatch_first(self) -> None: + """Test forward pass with batch first.""" + model, = self.node.f( + embed_dim=16, + num_heads=4, + dropout=0, + bias=True, + kdim=16, + vdim=16, + batch_first=True + ) + + self.assertIsInstance(model, PtnMultiheadAttentionCustomDef) + + # Create dummy input: batch_size=2, seq_len=5, embed_dim=16 + x = torch.rand(2, 5, 16) + attention_custom_mask = torch.ones(2, 5) + attn_output = model(x, attention_custom_mask) + self.assertEqual(attn_output.shape, (2, 5, 16)) + + def test_forward_seq_first(self) -> None: + """Test forward pass with seq first.""" + model, = self.node.f( + embed_dim=16, + num_heads=4, + dropout=0.1, + bias=True, + kdim=16, + vdim=16, + batch_first=False + ) + + self.assertIsInstance(model, PtnMultiheadAttentionCustomDef) + + # Create dummy input: batch_size=2, seq_len=5, embed_dim=16 + # Use seq first. + x = torch.rand(5, 2, 16) + attention_custom_mask = torch.ones(2, 5) + attn_output = model(x, attention_custom_mask) + self.assertEqual(attn_output.shape, (5, 2, 16)) + + +if __name__ == "__main__": + unittest.main()