diff --git a/__init__.py b/__init__.py index e9091a0..b0316d3 100644 --- a/__init__.py +++ b/__init__.py @@ -128,6 +128,8 @@ from .modules.pytorch_wrapper.ptn_conv_model import PtnConvModel from .modules.pytorch_wrapper.ptn_linear import PtnLinear from .modules.pytorch_wrapper.ptn_linear_model import PtnLinearModel from .modules.pytorch_wrapper.ptn_max_pool2d import PtnMaxPool2d +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_resnet_model import PtnResnetModel from .modules.pytorch_wrapper.pto_adam import PtoAdam from .modules.pytorch_wrapper.pto_adamw import PtoAdamW @@ -280,6 +282,8 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtnLinear": PtnLinear, "PtnLinearModel": PtnLinearModel, "PtnMaxPool2d": PtnMaxPool2d, + "PtnPreAddChannelAxis": PtnPreAddChannelAxis, + "PtnPreFlatten": PtnPreFlatten, "PtnResnetModel": PtnResnetModel, "PtoAdam": PtoAdam, "PtoAdamW": PtoAdamW, @@ -432,6 +436,8 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtnLinear": "Ptn Linear", "PtnLinearModel": "Ptn Linear Model", "PtnMaxPool2d": "Ptn Max Pool2d", + "PtnPreAddChannelAxis": "Ptn Pre Add Channel Axis", + "PtnPreFlatten": "Ptn Pre Flatten", "PtnResnetModel": "Ptn Resnet Model", "PtoAdam": "Pto Adam", "PtoAdamW": "Pto AdamW", diff --git a/docs/building_a_model_from_scratch.md b/docs/building_a_model_from_scratch.md new file mode 100644 index 0000000..e9ca1c4 --- /dev/null +++ b/docs/building_a_model_from_scratch.md @@ -0,0 +1,12 @@ +# Building a Model from Scratch + +ComfyUI Pt Wrapper provides composite models that can be easily customized by adjusting parameters. However, you can also build a model from scratch by sequentially combining basic models such as `Conv2D` and `Linear`. + +To chain two models, use the **Pt Chained Model** node. This node also supports specifying post-processing after the forward pass of the second model (e.g., activation functions). To define the post-processing step, connect a **Ptf** node (e.g., `Ptf ReLU`). + +In some cases, preprocessing may be required before feeding data into the model. For instance, you can use the **Pt Pre Add Channel Axis** node to add a channel dimension to the input tensor or the **Pt Pre Flatten** node to flatten convolutional output before passing it to a linear layer. + +The following example workflows demonstrate this approach: + +- `./examples/workflows/fashion_mnist_eval_conv_chained_model.json`: Constructs a multi-layer convolutional network from scratch. +- `./examples/workflows/fashion_mnist_train_chained_model.json`: Uses the trained model for evaluation. diff --git a/docs/model_training.md b/docs/model_training.md index c31ef37..adbc5b5 100644 --- a/docs/model_training.md +++ b/docs/model_training.md @@ -7,13 +7,12 @@ Currently, following archituctures are supported (node names are in parentheses) * Convolutional Network (Ptn Conv Model) * ResNet (Ptn Resnet Model) -In addition, a single layer Linear model is supported as an example to illustrate the use (see the bottom of this page on how to use it): -* Linear (Ptn Linear) - You can adjust the model configuration including output dimension on the UI. The system is designed to be modular. For example, to switch from a linear model to a convolutional model, you only need to replace the model node (see screenshot below). +You can also build a model from scratch by sequentially combining basic models such as `Conv2D` and `Linear` by using the **Pt Chained Model** node. See [Building a Model from Scratch](building_a_model_from_scratch.md) for details. + You can choose to: * Use your own images * Use a public dataset diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index 7ffd658..d817ba0 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -120,6 +120,8 @@ | [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 Max Pool](ptn_max_pool2d.md) | A maxpool layer. | +| [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 Resnet Model](ptn_resnet_model.md) | A Resnet model consisting of multiple Resnet layers. | ## PyTorch wrapper - Optimizer | Node | Description | diff --git a/docs/reference/ptn_conv2d.md b/docs/reference/ptn_conv2d.md index e5a0cd4..55d574c 100644 --- a/docs/reference/ptn_conv2d.md +++ b/docs/reference/ptn_conv2d.md @@ -25,7 +25,7 @@ t. | dilation | String | | groups | Int | | bias | Boolean | -| padding_mode | Zeros | +| padding_mode | | ## Output | Data type | diff --git a/docs/reference/ptn_pre_add_channel_axis.md b/docs/reference/ptn_pre_add_channel_axis.md new file mode 100644 index 0000000..d7abf33 --- /dev/null +++ b/docs/reference/ptn_pre_add_channel_axis.md @@ -0,0 +1,18 @@ +# Ptn Pre Add Channel Axis +Adds a channel axis after the batch axis if the input is rank 3 (bs, h, w) +before the forward pass by the specified model. + +## Input +| Name | Data type | +|---|---| +| model | Ptmodel | + +## 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_pre_flatten.md b/docs/reference/ptn_pre_flatten.md new file mode 100644 index 0000000..eddaeb1 --- /dev/null +++ b/docs/reference/ptn_pre_flatten.md @@ -0,0 +1,17 @@ +# Ptn Pre Flatten +Flattens the input tensor before processing the tensor in the specified model. + +## Input +| Name | Data type | +|---|---| +| model | Ptmodel | + +## 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/fashion_mnist_eval_conv_chained_model.json b/examples/workflows/fashion_mnist_eval_conv_chained_model.json new file mode 100644 index 0000000..1e9710e --- /dev/null +++ b/examples/workflows/fashion_mnist_eval_conv_chained_model.json @@ -0,0 +1 @@ +{"last_node_id":73,"last_link_id":88,"nodes":[{"id":47,"type":"PtEvaluateClassificationModel","pos":[3621.56201171875,4297.65869140625],"size":[403.1999816894531,138],"flags":{},"order":25,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":55},{"name":"data_loader","type":"PTDATALOADER","link":56}],"outputs":[{"name":"accuracy","type":"FLOAT","links":[57],"slot_index":0},{"name":"precision","type":"FLOAT","links":[58],"slot_index":1},{"name":"recall","type":"FLOAT","links":[59],"slot_index":2},{"name":"f1","type":"FLOAT","links":[60],"slot_index":3},{"name":"num_samples","type":"INT","links":[61],"slot_index":4}],"properties":{"Node name for S&R":"PtEvaluateClassificationModel"},"widgets_values":[true]},{"id":33,"type":"CDAShowInt","pos":[4094.39892578125,4483.74853515625],"size":[210,79.09933471679688],"flags":{},"order":30,"mode":0,"inputs":[{"name":"int_scalar","type":"INT","link":61,"widget":{"name":"int_scalar"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"title":"Num Samples","properties":{"Node name for S&R":"CDAShowInt"},"widgets_values":[0,"10000"]},{"id":31,"type":"CDAShowFloat","pos":[4315.8994140625,4342.650390625],"size":[210,76],"flags":{},"order":29,"mode":0,"inputs":[{"name":"float_scalar","type":"FLOAT","link":60,"widget":{"name":"float_scalar"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"title":"F1","properties":{"Node name for S&R":"CDAShowFloat"},"widgets_values":[0,"0.9051993263349519"]},{"id":29,"type":"CDAShowFloat","pos":[4315.8994140625,4218.80322265625],"size":[210,76],"flags":{},"order":27,"mode":0,"inputs":[{"name":"float_scalar","type":"FLOAT","link":58,"widget":{"name":"float_scalar"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"title":"Precision","properties":{"Node name for S&R":"CDAShowFloat"},"widgets_values":[0,"0.9074881608338703"]},{"id":30,"type":"CDAShowFloat","pos":[4094.39892578125,4342.650390625],"size":[210,76],"flags":{},"order":28,"mode":0,"inputs":[{"name":"float_scalar","type":"FLOAT","link":59,"widget":{"name":"float_scalar"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"title":"Recall","properties":{"Node name for S&R":"CDAShowFloat"},"widgets_values":[0,"0.9059000000000001"]},{"id":28,"type":"CDAShowFloat","pos":[4094.39892578125,4218.80322265625],"size":[210,76],"flags":{},"order":26,"mode":0,"inputs":[{"name":"float_scalar","type":"FLOAT","link":57,"widget":{"name":"float_scalar"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"title":"Accuracy","properties":{"Node name for S&R":"CDAShowFloat"},"widgets_values":[0,"0.9059"]},{"id":49,"type":"PtnMaxPool2d","pos":[2414.755615234375,1868.3812255859375],"size":[315,130],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[64],"slot_index":0}],"properties":{"Node name for S&R":"PtnMaxPool2d"},"widgets_values":["2","2","0","1"]},{"id":50,"type":"PtnChainedModel","pos":[2807.98876953125,1796.02197265625],"size":[214.20001220703125,66],"flags":{},"order":15,"mode":0,"inputs":[{"name":"model_a","type":"PTMODEL","link":63},{"name":"model_b","type":"PTMODEL","link":64},{"name":"closure","type":"PTCALLABLE","link":65,"shape":7}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[74],"slot_index":0}],"properties":{"Node name for S&R":"PtnChainedModel"},"widgets_values":[]},{"id":51,"type":"PtnConv2d","pos":[2414.755615234375,2190.901123046875],"size":[315,250],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[75],"slot_index":0}],"properties":{"Node name for S&R":"PtnConv2d"},"widgets_values":[32,64,"3","1","same","1",1,true,"zeros"]},{"id":52,"type":"PtnMaxPool2d","pos":[2414.755615234375,2530.73388671875],"size":[315,130],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[77],"slot_index":0}],"properties":{"Node name for S&R":"PtnMaxPool2d"},"widgets_values":["2","2","0","1"]},{"id":53,"type":"PtnChainedModel","pos":[2807.98876953125,2868.0244140625],"size":[214.20001220703125,66],"flags":{},"order":18,"mode":0,"inputs":[{"name":"model_a","type":"PTMODEL","link":66},{"name":"model_b","type":"PTMODEL","link":67},{"name":"closure","type":"PTCALLABLE","link":null,"shape":7}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[68],"slot_index":0}],"properties":{"Node name for S&R":"PtnChainedModel"},"widgets_values":[]},{"id":54,"type":"PtnChainedModel","pos":[2807.98876953125,3124.349609375],"size":[214.20001220703125,66],"flags":{},"order":19,"mode":0,"inputs":[{"name":"model_a","type":"PTMODEL","link":68},{"name":"model_b","type":"PTMODEL","link":69},{"name":"closure","type":"PTCALLABLE","link":70,"shape":7}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[71],"slot_index":0}],"properties":{"Node name for S&R":"PtnChainedModel"},"widgets_values":[]},{"id":55,"type":"PtnConv2d","pos":[2414.755615234375,3367.09765625],"size":[315,250],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[72],"slot_index":0}],"properties":{"Node name for S&R":"PtnConv2d"},"widgets_values":[128,256,"3","1","same","1",1,true,"zeros"]},{"id":56,"type":"PtnChainedModel","pos":[2807.98876953125,3374.619140625],"size":[214.20001220703125,66],"flags":{},"order":20,"mode":0,"inputs":[{"name":"model_a","type":"PTMODEL","link":71},{"name":"model_b","type":"PTMODEL","link":72},{"name":"closure","type":"PTCALLABLE","link":null,"shape":7}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[79],"slot_index":0}],"properties":{"Node name for S&R":"PtnChainedModel"},"widgets_values":[]},{"id":58,"type":"MarkdownNote","pos":[2415.82958984375,1415.5794677734375],"size":[210,58.94380569458008],"flags":{},"order":4,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Define model"],"color":"#432","bgcolor":"#653"},{"id":60,"type":"PtnConv2d","pos":[2414.755615234375,1555.407470703125],"size":[315,250],"flags":{},"order":5,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[73],"slot_index":0}],"properties":{"Node name for S&R":"PtnConv2d"},"widgets_values":[1,32,"3","1","same","1",1,true,"zeros"]},{"id":61,"type":"PtnPreAddChannelAxis","pos":[2807.98876953125,1554.4034423828125],"size":[211.08056640625,30.150875091552734],"flags":{},"order":13,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":73}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[63],"slot_index":0}],"properties":{"Node name for S&R":"PtnPreAddChannelAxis"},"widgets_values":[]},{"id":62,"type":"PtnChainedModel","pos":[2807.98876953125,2341.434326171875],"size":[214.20001220703125,66],"flags":{},"order":16,"mode":0,"inputs":[{"name":"model_a","type":"PTMODEL","link":74},{"name":"model_b","type":"PTMODEL","link":75},{"name":"closure","type":"PTCALLABLE","link":null,"shape":7}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[76],"slot_index":0}],"properties":{"Node name for S&R":"PtnChainedModel"},"widgets_values":[]},{"id":63,"type":"PtnChainedModel","pos":[2807.98876953125,2527.0712890625],"size":[214.20001220703125,66],"flags":{},"order":17,"mode":0,"inputs":[{"name":"model_a","type":"PTMODEL","link":76},{"name":"model_b","type":"PTMODEL","link":77},{"name":"closure","type":"PTCALLABLE","link":78,"shape":7}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[66],"slot_index":0}],"properties":{"Node name for S&R":"PtnChainedModel"},"widgets_values":[]},{"id":64,"type":"PtnMaxPool2d","pos":[2414.755615234375,3692.7568359375],"size":[315,130],"flags":{},"order":6,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[80],"slot_index":0}],"properties":{"Node name for S&R":"PtnMaxPool2d"},"widgets_values":["2","2","0","1"]},{"id":65,"type":"PtnConv2d","pos":[2414.755615234375,2791.902099609375],"size":[315,250],"flags":{},"order":7,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[67],"slot_index":0}],"properties":{"Node name for S&R":"PtnConv2d"},"widgets_values":[64,128,"3","1","same","1",1,true,"zeros"]},{"id":66,"type":"PtnMaxPool2d","pos":[2414.755615234375,3117.850341796875],"size":[315,130],"flags":{},"order":8,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[69],"slot_index":0}],"properties":{"Node name for S&R":"PtnMaxPool2d"},"widgets_values":["2","2","0","1"]},{"id":67,"type":"PtnChainedModel","pos":[2807.98876953125,3684.239990234375],"size":[214.20001220703125,66],"flags":{},"order":21,"mode":0,"inputs":[{"name":"model_a","type":"PTMODEL","link":79},{"name":"model_b","type":"PTMODEL","link":80},{"name":"closure","type":"PTCALLABLE","link":81,"shape":7}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[83],"slot_index":0}],"properties":{"Node name for S&R":"PtnChainedModel"},"widgets_values":[]},{"id":68,"type":"PtnConv2d","pos":[2414.755615234375,3954.343505859375],"size":[315,250],"flags":{},"order":9,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[84],"slot_index":0}],"properties":{"Node name for S&R":"PtnConv2d"},"widgets_values":[256,512,"1","1","0","1",1,false,"zeros"]},{"id":69,"type":"PtfReLU","pos":[1922.4840087890625,4000.2392578125],"size":[210,26],"flags":{},"order":10,"mode":0,"inputs":[],"outputs":[{"name":"PTCALLABLE","type":"PTCALLABLE","links":[65,70,78,81,85],"slot_index":0}],"properties":{"Node name for S&R":"PtfReLU"},"widgets_values":[]},{"id":70,"type":"PtnPreFlatten","pos":[2414.755615234375,4306.3447265625],"size":[210,26],"flags":{},"order":14,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":82}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[87],"slot_index":0}],"properties":{"Node name for S&R":"PtnPreFlatten"},"widgets_values":[]},{"id":71,"type":"PtnChainedModel","pos":[2807.98876953125,4033.70654296875],"size":[214.20001220703125,66],"flags":{},"order":22,"mode":0,"inputs":[{"name":"model_a","type":"PTMODEL","link":83},{"name":"model_b","type":"PTMODEL","link":84},{"name":"closure","type":"PTCALLABLE","link":85,"shape":7}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[86],"slot_index":0}],"properties":{"Node name for S&R":"PtnChainedModel"},"widgets_values":[]},{"id":73,"type":"PtnLinear","pos":[2032.3580322265625,4306.3447265625],"size":[315,106],"flags":{},"order":11,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[82],"slot_index":0}],"properties":{"Node name for S&R":"PtnLinear"},"widgets_values":[512,10,true]},{"id":46,"type":"PtvDatasetLoader","pos":[3262.39892578125,4319.14794921875],"size":[266.8697814941406,246],"flags":{},"order":12,"mode":0,"inputs":[],"outputs":[{"name":"PTDATALOADER","type":"PTDATALOADER","links":[56],"slot_index":0}],"properties":{"Node name for S&R":"PtvDatasetLoader"},"widgets_values":["FashionMNIST",true,"",128,false,"{\"train\":False}","{\"num_workers\":1}"]},{"id":72,"type":"PtnChainedModel","pos":[2807.98876953125,4306.3447265625],"size":[214.20001220703125,66],"flags":{},"order":23,"mode":0,"inputs":[{"name":"model_a","type":"PTMODEL","link":86},{"name":"model_b","type":"PTMODEL","link":87},{"name":"closure","type":"PTCALLABLE","link":null,"shape":7}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[88],"slot_index":0}],"properties":{"Node name for S&R":"PtnChainedModel"},"widgets_values":[]},{"id":23,"type":"PtLoadModel","pos":[3622.276123046875,4128.44677734375],"size":[397.5032653808594,67.37537384033203],"flags":{},"order":24,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":88}],"outputs":[{"name":"model","type":"PTMODEL","links":[55],"slot_index":0}],"properties":{"Node name for S&R":"PtLoadModel"},"widgets_values":["fashion_mnist_chained_model.pt"]}],"links":[[55,23,0,47,0,"PTMODEL"],[56,46,0,47,1,"PTDATALOADER"],[57,47,0,28,0,"FLOAT"],[58,47,1,29,0,"FLOAT"],[59,47,2,30,0,"FLOAT"],[60,47,3,31,0,"FLOAT"],[61,47,4,33,0,"INT"],[63,61,0,50,0,"PTMODEL"],[64,49,0,50,1,"PTMODEL"],[65,69,0,50,2,"PTCALLABLE"],[66,63,0,53,0,"PTMODEL"],[67,65,0,53,1,"PTMODEL"],[68,53,0,54,0,"PTMODEL"],[69,66,0,54,1,"PTMODEL"],[70,69,0,54,2,"PTCALLABLE"],[71,54,0,56,0,"PTMODEL"],[72,55,0,56,1,"PTMODEL"],[73,60,0,61,0,"PTMODEL"],[74,50,0,62,0,"PTMODEL"],[75,51,0,62,1,"PTMODEL"],[76,62,0,63,0,"PTMODEL"],[77,52,0,63,1,"PTMODEL"],[78,69,0,63,2,"PTCALLABLE"],[79,56,0,67,0,"PTMODEL"],[80,64,0,67,1,"PTMODEL"],[81,69,0,67,2,"PTCALLABLE"],[82,73,0,70,0,"PTMODEL"],[83,67,0,71,0,"PTMODEL"],[84,68,0,71,1,"PTMODEL"],[85,69,0,71,2,"PTCALLABLE"],[86,71,0,72,0,"PTMODEL"],[87,70,0,72,1,"PTMODEL"],[88,72,0,23,0,"PTMODEL"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.37462470643979134,"offset":[-2187.1833905865665,-2735.6165457972957]},"node_versions":{"pt-wrapper":"a9870974df89837b31ce0a68cd9cd91c1fd6eb6a","data-analysis":"e490b620286600930c8c9d3a6986f73d5320463f"},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/examples/workflows/fashion_mnist_train_chained_model.json b/examples/workflows/fashion_mnist_train_chained_model.json new file mode 100644 index 0000000..4cf2f0e --- /dev/null +++ b/examples/workflows/fashion_mnist_train_chained_model.json @@ -0,0 +1 @@ +{"last_node_id":87,"last_link_id":133,"nodes":[{"id":50,"type":"PtTrainClassificationModel","pos":[2205.1328125,2001.00634765625],"size":[243.70773315429688,214],"flags":{},"order":32,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":118,"localized_name":"model"},{"name":"train_loader","type":"PTDATALOADER","link":66,"localized_name":"train_loader"},{"name":"optimizer","type":"PTOPTIMIZER","link":68,"localized_name":"optimizer"},{"name":"val_loader","type":"PTDATALOADER","link":67,"shape":7,"localized_name":"val_loader"}],"outputs":[{"name":"Model","type":"PTMODEL","links":[71],"slot_index":0,"localized_name":"Model"},{"name":"train loss","type":"TENSOR","links":null,"localized_name":"train loss"},{"name":"val loss","type":"TENSOR","links":null,"localized_name":"val loss"}],"properties":{"Node name for S&R":"PtTrainClassificationModel"},"widgets_values":[5,true,false,10,true]},{"id":11,"type":"PtoAdam","pos":[2205.1328125,1709.9505615234375],"size":[265.8851013183594,108.32855224609375],"flags":{},"order":31,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":117,"localized_name":"model"}],"outputs":[{"name":"PTOPTIMIZER","type":"PTOPTIMIZER","links":[68],"slot_index":0,"localized_name":"PTOPTIMIZER"}],"properties":{"Node name for S&R":"PtoAdam"},"widgets_values":[0.001,0.9,0.999]},{"id":56,"type":"MarkdownNote","pos":[2502.79443359375,2070.234130859375],"size":[210,58.94380569458008],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Save model"],"color":"#432","bgcolor":"#653"},{"id":51,"type":"PandasShowText","pos":[2503.911376953125,2185.8603515625],"size":[576.9476318359375,81.15640258789062],"flags":{},"order":34,"mode":0,"inputs":[{"name":"text","type":"STRING","link":70,"widget":{"name":"text"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6,"localized_name":"STRING"}],"properties":{"Node name for S&R":"PandasShowText"},"widgets_values":["","/home/dockeruser/comfy/ComfyUI/custom_nodes/pt-wrapper/models/fashion_mnist_chained_model.pt"]},{"id":54,"type":"MarkdownNote","pos":[1491.0869140625,2225.249755859375],"size":[210,58.94380569458008],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Load val data"],"color":"#432","bgcolor":"#653"},{"id":48,"type":"PtvDatasetLoader","pos":[1804.6444091796875,2226.46533203125],"size":[309.6732482910156,246.00001525878906],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[{"name":"PTDATALOADER","type":"PTDATALOADER","links":[67],"slot_index":0,"localized_name":"PTDATALOADER"}],"properties":{"Node name for S&R":"PtvDatasetLoader"},"widgets_values":["FashionMNIST",true,"",128,false,"{\"train\":False}","{\"num_workers\":1}"]},{"id":47,"type":"PtvDatasetLoader","pos":[1804.6444091796875,1913.2149658203125],"size":[314.89654541015625,246],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[{"name":"PTDATALOADER","type":"PTDATALOADER","links":[66],"slot_index":0,"localized_name":"PTDATALOADER"}],"properties":{"Node name for S&R":"PtvDatasetLoader"},"widgets_values":["FashionMNIST",true,"",128,true,"{\"train\":True}","{\"num_workers\":1}"]},{"id":72,"type":"PtnMaxPool2d","pos":[626.22412109375,-96.74343872070312],"size":[315,130],"flags":{},"order":4,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[91],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnMaxPool2d"},"widgets_values":["2","2","0","1"]},{"id":73,"type":"PtnChainedModel","pos":[1019.4572143554688,-169.1027374267578],"size":[214.20001220703125,66],"flags":{},"order":22,"mode":0,"inputs":[{"name":"model_a","type":"PTMODEL","link":124,"localized_name":"model_a"},{"name":"model_b","type":"PTMODEL","link":91,"localized_name":"model_b"},{"name":"closure","type":"PTCALLABLE","link":93,"shape":7,"localized_name":"closure"}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[99],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnChainedModel"},"widgets_values":[]},{"id":62,"type":"PtnConv2d","pos":[626.22412109375,225.7763214111328],"size":[315,250],"flags":{},"order":5,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[100],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnConv2d"},"widgets_values":[32,64,"3","1","same","1",1,true,"zeros"]},{"id":74,"type":"PtnMaxPool2d","pos":[626.22412109375,565.6092529296875],"size":[315,130],"flags":{},"order":6,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[102],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnMaxPool2d"},"widgets_values":["2","2","0","1"]},{"id":79,"type":"PtnChainedModel","pos":[1019.4572143554688,902.8997192382812],"size":[214.20001220703125,66],"flags":{},"order":25,"mode":0,"inputs":[{"name":"model_a","type":"PTMODEL","link":103,"localized_name":"model_a"},{"name":"model_b","type":"PTMODEL","link":104,"localized_name":"model_b"},{"name":"closure","type":"PTCALLABLE","link":null,"shape":7,"localized_name":"closure"}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[105],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnChainedModel"},"widgets_values":[]},{"id":64,"type":"PtnChainedModel","pos":[1019.4572143554688,1159.224853515625],"size":[214.20001220703125,66],"flags":{},"order":26,"mode":0,"inputs":[{"name":"model_a","type":"PTMODEL","link":105,"localized_name":"model_a"},{"name":"model_b","type":"PTMODEL","link":106,"localized_name":"model_b"},{"name":"closure","type":"PTCALLABLE","link":76,"shape":7,"localized_name":"closure"}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[107],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnChainedModel"},"widgets_values":[]},{"id":58,"type":"PtnConv2d","pos":[626.22412109375,1401.9730224609375],"size":[315,250],"flags":{},"order":7,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[108],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnConv2d"},"widgets_values":[128,256,"3","1","same","1",1,true,"zeros"]},{"id":81,"type":"PtnChainedModel","pos":[1019.4572143554688,1409.4945068359375],"size":[214.20001220703125,66],"flags":{},"order":27,"mode":0,"inputs":[{"name":"model_a","type":"PTMODEL","link":107,"localized_name":"model_a"},{"name":"model_b","type":"PTMODEL","link":108,"localized_name":"model_b"},{"name":"closure","type":"PTCALLABLE","link":null,"shape":7,"localized_name":"closure"}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[109],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnChainedModel"},"widgets_values":[]},{"id":57,"type":"PtnConvModel","pos":[52.205570220947266,549.4094848632812],"size":[315,202],"flags":{},"order":8,"mode":2,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnConvModel"},"widgets_values":["(1,28,28)",0,10,"[32,64,128,256,512]","[3,3,3,3,1]","[1,1,1,1,0]","[True,True,True,True,False]"]},{"id":52,"type":"MarkdownNote","pos":[627.2979736328125,-549.5452880859375],"size":[210,58.94380569458008],"flags":{},"order":9,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Define model"],"color":"#432","bgcolor":"#653"},{"id":82,"type":"MarkdownNote","pos":[47.48865509033203,397.1159973144531],"size":[336.3369445800781,82.82988739013672],"flags":{},"order":10,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Equivalent of below model"],"color":"#432","bgcolor":"#653"},{"id":53,"type":"MarkdownNote","pos":[1491.0869140625,1916.8900146484375],"size":[210,58.94380569458008],"flags":{},"order":11,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Load train data"],"color":"#432","bgcolor":"#653"},{"id":55,"type":"MarkdownNote","pos":[2205.1328125,2272.39453125],"size":[210,58.94380569458008],"flags":{},"order":12,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Train"],"color":"#432","bgcolor":"#653"},{"id":19,"type":"PtSaveModel","pos":[2499.0068359375,1958.690673828125],"size":[418.03662109375,64.53935241699219],"flags":{},"order":33,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":71,"localized_name":"model"}],"outputs":[{"name":"model full path","type":"STRING","links":[70],"slot_index":0,"localized_name":"model full path"}],"properties":{"Node name for S&R":"PtSaveModel"},"widgets_values":["fashion_mnist_chained_model.pt"]},{"id":70,"type":"PtnConv2d","pos":[626.22412109375,-409.7171936035156],"size":[315,250],"flags":{},"order":13,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[121],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnConv2d"},"widgets_values":[1,32,"3","1","same","1",1,true,"zeros"]},{"id":84,"type":"PtnPreAddChannelAxis","pos":[1019.4572143554688,-410.7212829589844],"size":[211.08056640625,30.150875091552734],"flags":{},"order":20,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":121,"localized_name":"model"}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[124],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnPreAddChannelAxis"},"widgets_values":[]},{"id":78,"type":"PtnChainedModel","pos":[1019.4572143554688,376.3096008300781],"size":[214.20001220703125,66],"flags":{},"order":23,"mode":0,"inputs":[{"name":"model_a","type":"PTMODEL","link":99,"localized_name":"model_a"},{"name":"model_b","type":"PTMODEL","link":100,"localized_name":"model_b"},{"name":"closure","type":"PTCALLABLE","link":null,"shape":7,"localized_name":"closure"}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[125],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnChainedModel"},"widgets_values":[]},{"id":71,"type":"PtnChainedModel","pos":[1019.4572143554688,561.9466552734375],"size":[214.20001220703125,66],"flags":{},"order":24,"mode":0,"inputs":[{"name":"model_a","type":"PTMODEL","link":125,"localized_name":"model_a"},{"name":"model_b","type":"PTMODEL","link":102,"localized_name":"model_b"},{"name":"closure","type":"PTCALLABLE","link":89,"shape":7,"localized_name":"closure"}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[103],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnChainedModel"},"widgets_values":[]},{"id":76,"type":"PtnMaxPool2d","pos":[626.22412109375,1727.632080078125],"size":[315,130],"flags":{},"order":14,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[110],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnMaxPool2d"},"widgets_values":["2","2","0","1"]},{"id":61,"type":"PtnConv2d","pos":[626.22412109375,826.7774047851562],"size":[315,250],"flags":{},"order":15,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[104],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnConv2d"},"widgets_values":[64,128,"3","1","same","1",1,true,"zeros"]},{"id":75,"type":"PtnMaxPool2d","pos":[626.22412109375,1152.7255859375],"size":[315,130],"flags":{},"order":16,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[106],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnMaxPool2d"},"widgets_values":["2","2","0","1"]},{"id":67,"type":"PtnChainedModel","pos":[1019.4572143554688,1719.1153564453125],"size":[214.20001220703125,66],"flags":{},"order":28,"mode":0,"inputs":[{"name":"model_a","type":"PTMODEL","link":109,"localized_name":"model_a"},{"name":"model_b","type":"PTMODEL","link":110,"localized_name":"model_b"},{"name":"closure","type":"PTCALLABLE","link":79,"shape":7,"localized_name":"closure"}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[129],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnChainedModel"},"widgets_values":[]},{"id":59,"type":"PtnConv2d","pos":[626.22412109375,1989.2188720703125],"size":[315,250],"flags":{},"order":17,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[130],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnConv2d"},"widgets_values":[256,512,"1","1","0","1",1,false,"zeros"]},{"id":65,"type":"PtfReLU","pos":[133.952392578125,2035.114501953125],"size":[210,26],"flags":{},"order":18,"mode":0,"inputs":[],"outputs":[{"name":"PTCALLABLE","type":"PTCALLABLE","links":[76,79,89,93,131],"slot_index":0,"localized_name":"PTCALLABLE"}],"properties":{"Node name for S&R":"PtfReLU"},"widgets_values":[]},{"id":85,"type":"PtnPreFlatten","pos":[626.22412109375,2341.22021484375],"size":[210,26],"flags":{},"order":21,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":128,"localized_name":"model"}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[132],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnPreFlatten"},"widgets_values":[]},{"id":87,"type":"PtnChainedModel","pos":[1019.4572143554688,2068.581787109375],"size":[214.20001220703125,66],"flags":{},"order":29,"mode":0,"inputs":[{"name":"model_a","type":"PTMODEL","link":129,"localized_name":"model_a"},{"name":"model_b","type":"PTMODEL","link":130,"localized_name":"model_b"},{"name":"closure","type":"PTCALLABLE","link":131,"shape":7,"localized_name":"closure"}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[133],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnChainedModel"},"widgets_values":[]},{"id":80,"type":"PtnChainedModel","pos":[1019.4572143554688,2341.22021484375],"size":[214.20001220703125,66],"flags":{},"order":30,"mode":0,"inputs":[{"name":"model_a","type":"PTMODEL","link":133,"localized_name":"model_a"},{"name":"model_b","type":"PTMODEL","link":132,"localized_name":"model_b"},{"name":"closure","type":"PTCALLABLE","link":null,"shape":7,"localized_name":"closure"}],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[117,118],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnChainedModel"},"widgets_values":[]},{"id":86,"type":"PtnLinear","pos":[243.826416015625,2341.22021484375],"size":[315,106],"flags":{},"order":19,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[128],"slot_index":0,"localized_name":"PTMODEL"}],"properties":{"Node name for S&R":"PtnLinear"},"widgets_values":[512,10,true]}],"links":[[66,47,0,50,1,"PTDATALOADER"],[67,48,0,50,3,"PTDATALOADER"],[68,11,0,50,2,"PTOPTIMIZER"],[70,19,0,51,0,"STRING"],[71,50,0,19,0,"PTMODEL"],[76,65,0,64,2,"PTCALLABLE"],[79,65,0,67,2,"PTCALLABLE"],[89,65,0,71,2,"PTCALLABLE"],[91,72,0,73,1,"PTMODEL"],[93,65,0,73,2,"PTCALLABLE"],[99,73,0,78,0,"PTMODEL"],[100,62,0,78,1,"PTMODEL"],[102,74,0,71,1,"PTMODEL"],[103,71,0,79,0,"PTMODEL"],[104,61,0,79,1,"PTMODEL"],[105,79,0,64,0,"PTMODEL"],[106,75,0,64,1,"PTMODEL"],[107,64,0,81,0,"PTMODEL"],[108,58,0,81,1,"PTMODEL"],[109,81,0,67,0,"PTMODEL"],[110,76,0,67,1,"PTMODEL"],[117,80,0,11,0,"PTMODEL"],[118,80,0,50,0,"PTMODEL"],[121,70,0,84,0,"PTMODEL"],[124,84,0,73,0,"PTMODEL"],[125,78,0,71,0,"PTMODEL"],[128,86,0,85,0,"PTMODEL"],[129,67,0,87,0,"PTMODEL"],[130,59,0,87,1,"PTMODEL"],[131,65,0,87,2,"PTCALLABLE"],[132,85,0,80,1,"PTMODEL"],[133,87,0,80,0,"PTMODEL"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.6033368359683484,"offset":[-1953.2206104686907,-3068.2635967230185]},"node_versions":{"pt-wrapper":"a9870974df89837b31ce0a68cd9cd91c1fd6eb6a","data-analysis":"e490b620286600930c8c9d3a6986f73d5320463f"},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/modules/pytorch_wrapper/ptn_conv_model_def.py b/modules/pytorch_wrapper/ptn_conv_model_def.py index e453c9c..7f23bef 100644 --- a/modules/pytorch_wrapper/ptn_conv_model_def.py +++ b/modules/pytorch_wrapper/ptn_conv_model_def.py @@ -29,12 +29,13 @@ class ConvSingleLayer(nn.Module): kernel_size: int, padding: int, downsample: bool, - bias:bool = False, + bias:bool = True, apply_activation: bool = True) -> None: super().__init__() self.conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, + bias=bias, padding=padding) self.apply_activation = apply_activation self.activation = nn.ReLU() diff --git a/modules/pytorch_wrapper/ptn_pre_add_channel_axis.py b/modules/pytorch_wrapper/ptn_pre_add_channel_axis.py new file mode 100644 index 0000000..bc49170 --- /dev/null +++ b/modules/pytorch_wrapper/ptn_pre_add_channel_axis.py @@ -0,0 +1,84 @@ +from typing import Any, Dict, Tuple, Callable, Optional +import torch +import torch.nn as nn + + +class PtnPreAddChannelAxisDef(nn.Module): + """ + Adds a channel axis after the batch axis if the input is rank 3 (bs, h, w). + + pragma: skip_doc + """ + + def __init__(self, + model: nn.Module + ): + super().__init__() + self.model = model + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + """ + Forward pass of the model. + + Args: + inputs (torch.Tensor): Input tensor. + + Returns: + torch.Tensor: Output tensor of the `model`. + """ + if inputs.dim() not in [3, 4]: + raise ValueError("Inputs is not a rank 3 or 4 tensor.") + + if inputs.dim() == 3: + x = torch.unsqueeze(inputs, 1) + else: + x = inputs + + x = self.model(x) + return x + + +class PtnPreAddChannelAxis: + """ + Adds a channel axis after the batch axis if the input is rank 3 (bs, h, w) + before the forward pass by the specified model. + + 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", {}) + } + } + + 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): + """ + Adds a channel axis after the batch axis if the input is rank 3 (bs, h, w). + + Args: + model (nn.Module): The model. + + Returns: + Tuple[nn.Module]: A tuple containing the instantiated model. + """ + with torch.inference_mode(False): + output_model = PtnPreAddChannelAxisDef(model) + return (output_model,) diff --git a/modules/pytorch_wrapper/ptn_pre_flatten.py b/modules/pytorch_wrapper/ptn_pre_flatten.py new file mode 100644 index 0000000..711c2d4 --- /dev/null +++ b/modules/pytorch_wrapper/ptn_pre_flatten.py @@ -0,0 +1,76 @@ +from typing import Any, Dict, Tuple +import torch +import torch.nn as nn + + +class PtnPreFlattenDef(nn.Module): + """ + Flattens the input tensor before processing the tensor in the specified model. + + pragma: skip_doc + """ + + def __init__(self, + model: nn.Module + ): + super().__init__() + self.model = model + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + """ + Forward pass of the model. + + Args: + inputs (torch.Tensor): Input tensor. + + Returns: + torch.Tensor: Output tensor of the `model`. + """ + x = torch.flatten(inputs, start_dim=1) # Flatten except batch-axis + x = self.model(x) + return x + + +class PtnPreFlatten: + """ + Flattens the input tensor before processing the tensor in the specified model. + + 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", {}) + } + } + + 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): + """ + Adds a channel axis after the batch axis if the input is rank 3 (bs, h, w). + + Args: + model (nn.Module): The model. + + Returns: + Tuple[nn.Module]: A tuple containing the instantiated model. + """ + with torch.inference_mode(False): + output_model = PtnPreFlattenDef(model) + return (output_model,) diff --git a/pyproject.toml b/pyproject.toml index 2802d5b..dc44806 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.0.12" +version = "3.0.13" license = {file = "LICENSE"} [project.urls] diff --git a/tests/pytorch_wrapper/ptn_pre_add_channel_axis_test.py b/tests/pytorch_wrapper/ptn_pre_add_channel_axis_test.py new file mode 100644 index 0000000..93374fc --- /dev/null +++ b/tests/pytorch_wrapper/ptn_pre_add_channel_axis_test.py @@ -0,0 +1,49 @@ +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_pre_add_channel_axis import PtnPreAddChannelAxis + +class DummyModel(nn.Module): + def forward(self, x): + return x # Pass-through model for testing + + +class TestPtnPreAddChannelAxis(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.model_node = PtnPreAddChannelAxis() + self.model = self.model_node.f(model=DummyModel())[0] + + + def test_valid_shapes(self): + """Test valid input shapes (rank 3 and 4).""" + rank3_tensor = torch.randn(2, 8, 8) # (bs, h, w) + rank4_tensor = torch.randn(2, 3, 8, 8) # (bs, c, h, w) + + result3 = self.model(rank3_tensor) + result4 = self.model(rank4_tensor) + + self.assertEqual(result3.shape, (2, 1, 8, 8), f"Unexpected shape for rank 3: {result3.shape}") + self.assertEqual(result4.shape, (2, 3, 8, 8), f"Unexpected shape for rank 4: {result4.shape}") + + def test_invalid_shapes(self): + """Test invalid input shapes (rank 2 and rank 5).""" + rank2_tensor = torch.randn(8, 8) # (h, w) + rank5_tensor = torch.randn(2, 3, 8, 8, 8) # (bs, c, d, h, w) + + with self.assertRaises(ValueError, msg="Rank 2 tensor did not raise ValueError"): + self.model(rank2_tensor) + + with self.assertRaises(ValueError, msg="Rank 5 tensor did not raise ValueError"): + self.model(rank5_tensor) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptn_pre_flatten_test.py b/tests/pytorch_wrapper/ptn_pre_flatten_test.py new file mode 100644 index 0000000..d347c08 --- /dev/null +++ b/tests/pytorch_wrapper/ptn_pre_flatten_test.py @@ -0,0 +1,40 @@ +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_pre_flatten import PtnPreFlatten + +class DummyModel(nn.Module): + def forward(self, x): + return x # Pass-through model for testing + + +class TestPtnPreFlatten(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.model_node = PtnPreFlatten() + self.model = self.model_node.f(model=DummyModel())[0] + + def test_valid_shapes(self): + """Test flattening on valid input shapes.""" + rank3_tensor = torch.randn(2, 8, 8) # (bs, h, w) + rank4_tensor = torch.randn(2, 3, 8, 8) # (bs, c, h, w) + rank5_tensor = torch.randn(2, 3, 4, 8, 8) # (bs, c, d, h, w) + + result3 = self.model(rank3_tensor) + result4 = self.model(rank4_tensor) + result5 = self.model(rank5_tensor) + + self.assertEqual(result3.shape, (2, 64), f"Unexpected shape for rank 3: {result3.shape}") + self.assertEqual(result4.shape, (2, 192), f"Unexpected shape for rank 4: {result4.shape}") + self.assertEqual(result5.shape, (2, 768), f"Unexpected shape for rank 5: {result5.shape}") + + +if __name__ == "__main__": + unittest.main()