diff --git a/README.md b/README.md index 0e47476..6fb626b 100644 --- a/README.md +++ b/README.md @@ -1,24 +1,42 @@ # ComfyUI-Data-Analysis -## Update (March 3, 2025) +## Update (March 5, 2025) Training workflow trains the FashionMNIST model from scratch with a single click of a button. Experimental PyTorch model training has been added (documentation coming soon). No coding required—training is as simple as a single click! -Try it out: +It is designed to be modular. For example, to switch from a linear model to a conv model, you can just replace the model node (see screenshot below). - Training: examples/fashion_mnist_train.json - Evaluation: examples/fashion_mnist_eval.json +Or you can just drag and drop below workflows to ComfyUI to check it out. Model is downloaded automatically: -This workflow trains a FashionMNIST model from scratch effortlessly. +Linear model +* Training: examples/fashion_mnist_train.json +* Evaluation: examples/fashion_mnist_eval.json -**Training workflow** +Convolutional model +* Training: examples/fashion_mnist_train_conv.json +* Evaluation: examples/fashion_mnist_eval_conv.json + +These training workflows train a FashionMNIST model from scratch effortlessly. + +### Training workflow +**Linear model** ![Train](docs/images/fashion_mnist_train.png) -**Eval workflow** +**Conv model** +![Train](docs/images/conv_train.png) + +*Close up of the conv node* +![Train](docs/images/conv_train2.png) + +### Eval workflow +**Linear model** ![Eval](docs/images/fashion_mnist_eval.png) -Note that the training feature is a proof of concept and currently supports only the dense model. Additional models may be added based on demand. +**Conv model** +![Eval](docs/images/conv_eval.png) + +Note that the training feature is a proof of concept and currently supports only the linear and conv models. Additional models may be added based on demand.
diff --git a/__init__.py b/__init__.py index d070694..baeddf4 100644 --- a/__init__.py +++ b/__init__.py @@ -214,6 +214,7 @@ from .modules.pytorch_wrapper.pt_var import PtVar from .modules.pytorch_wrapper.pt_view import PtView from .modules.pytorch_wrapper.pt_where import PtWhere from .modules.pytorch_wrapper.pt_zeros import PtZeros +from .modules.pytorch_wrapper.ptn_conv_model import PtnConvModel from .modules.pytorch_wrapper.ptn_linear_model import PtnLinearModel from .modules.pytorch_wrapper.pto_adam import PtoAdam from .modules.pytorch_wrapper.ptv_dataset import PtvDataset @@ -446,6 +447,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtView": PtView, "PtWhere": PtWhere, "PtZeros": PtZeros, + "PtnConvModel": PtnConvModel, "PtnLinearModel": PtnLinearModel, "PtoAdam": PtoAdam, "PtvDataset": PtvDataset, @@ -681,6 +683,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtView": "Pt View", "PtWhere": "Pt Where", "PtZeros": "Pt Zeros", + "PtnConvModel": "Ptn Conv Model", "PtnLinearModel": "Ptn Linear Model", "PtoAdam": "Pto Adam", "PtvDataset": "Ptv Dataset", diff --git a/docs/images/conv_eval.png b/docs/images/conv_eval.png new file mode 100644 index 0000000..03ae6b9 Binary files /dev/null and b/docs/images/conv_eval.png differ diff --git a/docs/images/conv_train.png b/docs/images/conv_train.png new file mode 100644 index 0000000..442c4f3 Binary files /dev/null and b/docs/images/conv_train.png differ diff --git a/docs/images/conv_train2.png b/docs/images/conv_train2.png new file mode 100644 index 0000000..de5a9a2 Binary files /dev/null and b/docs/images/conv_train2.png differ diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index c6e6374..1623b38 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -122,10 +122,6 @@ | [Pandas Log](pandas_log.md) | Applying the natural logarithm function to a pandas DataFrame and converting non-numeric values to NaN. | | [Pandas Sin](pandas_sin.md) | Applies the sine function to a pandas DataFrame and converting non-numeric values to NaN. | | [Pandas Tan](pandas_tan.md) | Applies the tangent function to a pandas DataFrame and converting non-numeric values to NaN. | -## Miscellaneous -| Node | Description | -| --- | --- | -| [Pt Flatten](pt_flatten.md) | Flattens a PyTorch tensor into a 1D tensor. | ## Numpy | Node | Description | | --- | --- | @@ -266,6 +262,7 @@ | [Pt Evaluate Classification Model](pt_evaluate_classification_model.md) | Performs inference on test data and computes evaluation metrics. | | [Pt Load Model](pt_load_model.md) | A wrapper class for saving a PyTorch model. | | [Pt Train Classification Model](pt_train_classification_model.md) | Trains a classification model using a given dataset, optimizer, and number of epochs. | +| [Ptn Conv Model](ptn_conv_model.md) | A convolutional model consisting of multiple convolutional layers. | | [Ptn Linear Model](ptn_linear_model.md) | A linear model consisting of dense layers. | | [Ptv Dataset](ptv_dataset.md) | A Torchvision Dataset class wrapper. | | [Ptv Dataset Loader](ptv_dataset_loader.md) | A node to combine the dataset and data loader into a single node. | @@ -273,6 +270,7 @@ ## PyTorch wrapper - Transform | Node | Description | | --- | --- | +| [Pt Flatten](pt_flatten.md) | Flattens a PyTorch tensor into a 1D tensor. | | [Pt Permute](pt_permute.md) | Permutes the dimensions of a PyTorch tensor according to the specified order. | | [Pt Reshape](pt_reshape.md) | Reshapes a PyTorch tensor into a specified shape using `torch.reshape()`. | | [Pt Squeeze](pt_squeeze.md) | Removes a dimension at the specified position in the input tensor if it is of size 1. | diff --git a/docs/reference/pt_flatten.md b/docs/reference/pt_flatten.md index c9f79d6..9c32483 100644 --- a/docs/reference/pt_flatten.md +++ b/docs/reference/pt_flatten.md @@ -1,8 +1,6 @@ # Pt Flatten Flattens a PyTorch tensor into a 1D tensor. -Category: PyTorch wrapper - ## Input | Name | Data type | |---|---| @@ -14,6 +12,6 @@ Category: PyTorch wrapper | Tensor |
-Category: Miscellaneous +Category: PyTorch wrapper - Transform ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptn_conv_model.md b/docs/reference/ptn_conv_model.md new file mode 100644 index 0000000..d881ea0 --- /dev/null +++ b/docs/reference/ptn_conv_model.md @@ -0,0 +1,45 @@ +# Ptn Conv Model +A convolutional model consisting of multiple convolutional layers. + +### Fields: +- `input_dim`: A string representing the input dimensions in the format "(C,H,W)". + Example: `"(3,28,28)"` for a 3-channel 28x28 image. + +- `penultimate_dim`: An integer specifying the number of features before the output layer. If you specify 0, the number is computed internally. + Default: `0` + +- `output_dim`: An integer specifying the number of output features, which should match the number of target classes in classification. + Default: `10`, Min: `1`, Max: `1e6`. + +- `channel_list`: A string representing a Python list of integers specifying the number of channels per layer excluding the channel for the input (e.g. 3 for the color image). + Example: `"[32,64,128,256,512]"`. + +- `kernel_size_list`: A string representing a Python list of kernel sizes per layer. + Example: `"[3,3,3,3,1]"`. + +- `padding_list`: A string representing a Python list of padding values per layer. + Example: `"[1,1,1,1,0]"`. + +- `downsample_list`: A string representing a Python list of boolean values indicating whether to downsample at each layer. + Example: `"[True,True,True,True,False]"`. + +## Input +| Name | Data type | +|---|---| +| input_dim | String | +| penultimate_dim | Int | +| output_dim | Int | +| channel_list | String | +| kernel_size_list | String | +| padding_list | String | +| downsample_list | String | + +## Output +| Data type | +|---| +| Ptmodel | + +
+Category: PyTorch wrapper - Training + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptn_conv_model_def.md b/docs/reference/ptn_conv_model_def.md new file mode 100644 index 0000000..3a253f6 --- /dev/null +++ b/docs/reference/ptn_conv_model_def.md @@ -0,0 +1,17 @@ +# Conv Single Layer +A single convolutional layer followed by an optional ReLU activation and optional downsampling. + +Shape is tensor of shape `(batch_size, in_channels, height, width)`. + +Args: + in_channels (int): Number of input channels. + out_channels (int): Number of output channels. + kernel_size (int): Size of the convolutional kernel. + padding (int): Padding size for convolution. + downsample (bool): If True, applies a 2x2 MaxPooling layer. + apply_activation (bool, optional): If True, applies ReLU activation. Defaults to True. + +
+Category: Miscellaneous + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptn_linear_model.md b/docs/reference/ptn_linear_model.md index a13e46c..a9788fa 100644 --- a/docs/reference/ptn_linear_model.md +++ b/docs/reference/ptn_linear_model.md @@ -2,7 +2,7 @@ A linear model consisting of dense layers. ### Fields: -- `dim_list`: A comma-separated string representing layer dimensions. +- `dim_list`: A string representing a Python list of layer dimensions. It should contain the input dimension for the first layer and the output dimensions for all layers. @@ -10,9 +10,10 @@ A linear model consisting of dense layers. - Input data: 784 - Layer 1: 784 → 128 - Layer 2: 128 → 10 - - Specify `"784, 128, 10"` + - Specify `"[784,128,10]"` -- `bias_list`: A comma-separated string representing bias usage (True/False) per layer. +- `bias_list`: A string representing a Python list of bias usage (True/False) per layer. + Example: "[True,False]" - `num_layers`: The number of layers in the model. diff --git a/examples/workflows/fashion_mnist_eval_conv.json b/examples/workflows/fashion_mnist_eval_conv.json new file mode 100644 index 0000000..fc3883d --- /dev/null +++ b/examples/workflows/fashion_mnist_eval_conv.json @@ -0,0 +1 @@ +{"last_node_id":48,"last_link_id":62,"nodes":[{"id":47,"type":"PtEvaluateClassificationModel","pos":[3243.2138671875,1887.3746337890625],"size":[403.1999816894531,138],"flags":{},"order":3,"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":[3716.05078125,2073.464111328125],"size":[210,79.09933471679688],"flags":{},"order":8,"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":[3937.552490234375,1932.3660888671875],"size":[210,76],"flags":{},"order":7,"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.9085886169153918"]},{"id":29,"type":"CDAShowFloat","pos":[3937.552490234375,1808.5174560546875],"size":[210,76],"flags":{},"order":5,"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.9113484312788449"]},{"id":30,"type":"CDAShowFloat","pos":[3716.05078125,1932.3660888671875],"size":[210,76],"flags":{},"order":6,"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.908"]},{"id":28,"type":"CDAShowFloat","pos":[3716.05078125,1808.5174560546875],"size":[210,76],"flags":{},"order":4,"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.908"]},{"id":46,"type":"PtvDatasetLoader","pos":[2884.05078125,1908.8636474609375],"size":[266.8697814941406,246],"flags":{},"order":0,"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":48,"type":"PtnConvModel","pos":[2846.147216796875,1629.3045654296875],"size":[315,202],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[62],"slot_index":0}],"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":23,"type":"PtLoadModel","pos":[3243.927978515625,1718.1610107421875],"size":[397.5032653808594,67.37537384033203],"flags":{},"order":2,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":62}],"outputs":[{"name":"model","type":"PTMODEL","links":[55],"slot_index":0}],"properties":{"Node name for S&R":"PtLoadModel"},"widgets_values":["fashion_mnist_conv_ep5.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"],[62,48,0,23,0,"PTMODEL"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.9652303192711613,"offset":[-2706.3348752184625,-1469.581964047439]},"node_versions":{"data-analysis":"11335e88b00ed3b4a5129f92d3d134854b6b4a92"},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/examples/workflows/fashion_mnist_train_conv.json b/examples/workflows/fashion_mnist_train_conv.json new file mode 100644 index 0000000..c34d729 --- /dev/null +++ b/examples/workflows/fashion_mnist_train_conv.json @@ -0,0 +1 @@ +{"last_node_id":57,"last_link_id":73,"nodes":[{"id":50,"type":"PtTrainClassificationModel","pos":[3170.615478515625,1736.97705078125],"size":[243.70773315429688,142],"flags":{},"order":9,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":73},{"name":"train_loader","type":"PTDATALOADER","link":66},{"name":"optimizer","type":"PTOPTIMIZER","link":68},{"name":"val_loader","type":"PTDATALOADER","link":67,"shape":7}],"outputs":[{"name":"Model","type":"PTMODEL","links":[71],"slot_index":0},{"name":"train loss","type":"TENSOR","links":null},{"name":"val loss","type":"TENSOR","links":null}],"properties":{"Node name for S&R":"PtTrainClassificationModel"},"widgets_values":[5,true]},{"id":11,"type":"PtoAdam","pos":[3170.615478515625,1445.9212646484375],"size":[265.8851013183594,108.32855224609375],"flags":{},"order":8,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":72}],"outputs":[{"name":"PTOPTIMIZER","type":"PTOPTIMIZER","links":[68],"slot_index":0}],"properties":{"Node name for S&R":"PtoAdam"},"widgets_values":[0.001,0.9,0.999]},{"id":55,"type":"MarkdownNote","pos":[3170.615478515625,1953.1151123046875],"size":[210,58.94380569458008],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Train"],"color":"#432","bgcolor":"#653"},{"id":56,"type":"MarkdownNote","pos":[3468.277099609375,1806.20458984375],"size":[210,58.94380569458008],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Save model"],"color":"#432","bgcolor":"#653"},{"id":51,"type":"PandasShowText","pos":[3469.39404296875,1921.830078125],"size":[576.9476318359375,81.15640258789062],"flags":{},"order":11,"mode":0,"inputs":[{"name":"text","type":"STRING","link":70,"widget":{"name":"text"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"properties":{"Node name for S&R":"PandasShowText"},"widgets_values":["","/home/dockeruser/comfy/ComfyUI/custom_nodes/data-analysis/models/fashion_mnist_conv_ep5.pt"]},{"id":53,"type":"MarkdownNote","pos":[2456.569091796875,1652.8607177734375],"size":[210,58.94380569458008],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Load train data"],"color":"#432","bgcolor":"#653"},{"id":54,"type":"MarkdownNote","pos":[2456.569091796875,1961.218994140625],"size":[210,58.94380569458008],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Load val data"],"color":"#432","bgcolor":"#653"},{"id":52,"type":"MarkdownNote","pos":[2456.569091796875,1384.5592041015625],"size":[210,58.94380569458008],"flags":{},"order":4,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Define model"],"color":"#432","bgcolor":"#653"},{"id":48,"type":"PtvDatasetLoader","pos":[2770.126708984375,1962.4345703125],"size":[309.6732482910156,246.00001525878906],"flags":{},"order":5,"mode":0,"inputs":[],"outputs":[{"name":"PTDATALOADER","type":"PTDATALOADER","links":[67],"slot_index":0}],"properties":{"Node name for S&R":"PtvDatasetLoader"},"widgets_values":["FashionMNIST",true,"",128,false,"{\"train\":False}","{\"num_workers\":1}"]},{"id":47,"type":"PtvDatasetLoader","pos":[2770.126708984375,1649.1856689453125],"size":[314.89654541015625,246],"flags":{},"order":6,"mode":0,"inputs":[],"outputs":[{"name":"PTDATALOADER","type":"PTDATALOADER","links":[66],"slot_index":0}],"properties":{"Node name for S&R":"PtvDatasetLoader"},"widgets_values":["FashionMNIST",true,"",128,true,"{\"train\":True}","{\"num_workers\":1}"]},{"id":19,"type":"PtSaveModel","pos":[3464.489501953125,1694.661376953125],"size":[280.30267333984375,68.31287384033203],"flags":{},"order":10,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":71}],"outputs":[{"name":"model full path","type":"STRING","links":[70],"slot_index":0}],"properties":{"Node name for S&R":"PtSaveModel"},"widgets_values":["fashion_mnist_conv_ep5.pt"]},{"id":57,"type":"PtnConvModel","pos":[2770.126708984375,1380.5203857421875],"size":[315,202],"flags":{},"order":7,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[72,73],"slot_index":0}],"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]"]}],"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"],[72,57,0,11,0,"PTMODEL"],[73,57,0,50,0,"PTMODEL"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.7251918251473778,"offset":[-2209.359108864266,-1203.1203019096897]},"node_versions":{"data-analysis":"11335e88b00ed3b4a5129f92d3d134854b6b4a92"},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/modules/pytorch_wrapper/pt_flatten.py b/modules/pytorch_wrapper/pt_flatten.py index 4b565a2..ac56d3a 100644 --- a/modules/pytorch_wrapper/pt_flatten.py +++ b/modules/pytorch_wrapper/pt_flatten.py @@ -6,7 +6,7 @@ class PtFlatten: """ Flattens a PyTorch tensor into a 1D tensor. - Category: PyTorch wrapper + category: PyTorch wrapper - Transform """ @classmethod diff --git a/modules/pytorch_wrapper/ptn_conv_model.py b/modules/pytorch_wrapper/ptn_conv_model.py new file mode 100644 index 0000000..a31fe36 --- /dev/null +++ b/modules/pytorch_wrapper/ptn_conv_model.py @@ -0,0 +1,99 @@ +from typing import Any, Dict, Tuple +import ast +import torch +from .ptn_conv_model_def import ConvModel + + +class PtnConvModel: + """ + A convolutional model consisting of multiple convolutional layers. + + ### Fields: + - `input_dim`: A string representing the input dimensions in the format "(C,H,W)". + Example: `"(3,28,28)"` for a 3-channel 28x28 image. + + - `penultimate_dim`: An integer specifying the number of features before the output layer. If you specify 0, the number is computed internally. + Default: `0` + + - `output_dim`: An integer specifying the number of output features, which should match the number of target classes in classification. + Default: `10`, Min: `1`, Max: `1e6`. + + - `channel_list`: A string representing a Python list of integers specifying the number of channels per layer excluding the channel for the input (e.g. 3 for the color image). + Example: `"[32,64,128,256,512]"`. + + - `kernel_size_list`: A string representing a Python list of kernel sizes per layer. + Example: `"[3,3,3,3,1]"`. + + - `padding_list`: A string representing a Python list of padding values per layer. + Example: `"[1,1,1,1,0]"`. + + - `downsample_list`: A string representing a Python list of boolean values indicating whether to downsample at each layer. + Example: `"[True,True,True,True,False]"`. + + category: PyTorch wrapper - Training + """ + + @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": { + "input_dim": ("STRING", {"default": "(3,28,28)"}), + "penultimate_dim": ("INT", {"default": 0, "min": 0, "max": 1e6}), + "output_dim": ("INT", {"default": 10, "min": 1, "max": 1e6}), + "channel_list": ("STRING", {"default": "[32,64,128,256,512]"}), + "kernel_size_list": ("STRING", {"default": "[3,3,3,3,1]"}), + "padding_list": ("STRING", {"default": "[1,1,1,1,0]"}), + "downsample_list": ("STRING", {"default": "[True,True,True,True,False]"}) + } + } + + RETURN_TYPES: Tuple[str] = ("PTMODEL",) + FUNCTION: str = "f" + CATEGORY: str = "Training" + + @classmethod + def IS_CHANGED(cls, **kw): + return float("NaN") + + def f(self, + input_dim:str, + penultimate_dim: int, + output_dim: int, + channel_list: str, + kernel_size_list: str, + padding_list: str, + downsample_list: str) -> None: + """ + Constructs a convolutional neural network model. + + Args: + TODO:Update + + Returns: + Tuple[nn.Module]: A tuple containing the instantiated PyTorch model. + + Raises: + TODO:Update + """ + input_dim_parsed = tuple(ast.literal_eval(input_dim)) + channel_list_parsed = ast.literal_eval(channel_list) + kernel_size_list_parsed = ast.literal_eval(kernel_size_list) + padding_list_parsed = ast.literal_eval(padding_list) + downsample_list_parsed = ast.literal_eval(downsample_list) + + with torch.inference_mode(False): + model = ConvModel( + input_dim_parsed, + penultimate_dim, + output_dim, + channel_list_parsed, + kernel_size_list_parsed, + padding_list_parsed, + downsample_list_parsed) + return (model,) diff --git a/modules/pytorch_wrapper/ptn_conv_model_def.py b/modules/pytorch_wrapper/ptn_conv_model_def.py new file mode 100644 index 0000000..74f71c6 --- /dev/null +++ b/modules/pytorch_wrapper/ptn_conv_model_def.py @@ -0,0 +1,151 @@ +from typing import List, Dict, Any, Tuple +import torch +import torch.nn as nn + + +class ConvSingleLayer(nn.Module): + """A single convolutional layer followed by an optional ReLU activation and optional downsampling. + + Shape is tensor of shape `(batch_size, in_channels, height, width)`. + + Args: + in_channels (int): Number of input channels. + out_channels (int): Number of output channels. + kernel_size (int): Size of the convolutional kernel. + padding (int): Padding size for convolution. + downsample (bool): If True, applies a 2x2 MaxPooling layer. + apply_activation (bool, optional): If True, applies ReLU activation. Defaults to True. + + pragma: skip_doc + """ + + @classmethod # TODO: Dummy method to suppress ComfyUI warning. DO NOT REMOVE + def INPUT_TYPES(cls) -> Dict[str, Any]: + return {} + + def __init__(self, + in_channels: int, + out_channels: int, + kernel_size: int, + padding: int, + downsample: bool, + apply_activation: bool = True) -> None: + super().__init__() + self.conv = nn.Conv2d(in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + padding=padding) + self.apply_activation = apply_activation + self.activation = nn.ReLU() + self.downsample = downsample + self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2) + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + """Forward pass of the convolutional layer. + + Args: + inputs (torch.Tensor): Input tensor of shape `(batch_size, in_channels, height, width)`. + + Returns: + torch.Tensor: Output tensor after convolution, optional activation, and optional pooling. + """ + x = self.conv(inputs) + if self.downsample: + x = self.maxpool(x) + if self.apply_activation: + x = self.activation(x) + return x + + +class ConvModel(nn.Module): + """A multi-layer convolutional neural network with ReLU activation after each convolutional layer, except for the final dense layer. + + Args: + input_dim (Tuple[int, int, int]): Shape of a single input sample (C, H, W). + penultimate_dim (int): Number of features before the final dense layer. If 0 is specified, this value is automatically computed internally. + output_dim (int): Number of output features. + channel_list (List[int]): List of channel sizes for each layer. This should not include the inputs channel. + Input channel is taken from the first element of input_dim. + kernel_size_list (List[int]): List of kernel sizes for each convolutional layer. + padding_list (List[int]): List of padding values for each convolutional layer. + downsample_list (List[bool]): List indicating whether downsampling should be applied after each layer. + + pragma: skip_doc + """ + + @classmethod # TODO: Dummy method to suppress ComfyUI warning + def INPUT_TYPES(cls) -> Dict[str, Any]: + return {} + + def __init__(self, + input_dim: Tuple[int, int, int], + penultimate_dim: int, + output_dim: int, + channel_list: List[int], + kernel_size_list: List[int], + padding_list: List[int], + downsample_list: List[bool]) -> None: + super().__init__() + + # Validate list lengths + if (len(channel_list) != len(kernel_size_list) or + len(kernel_size_list) != len(padding_list) or + len(padding_list) != len(downsample_list)): + raise ValueError("channel_list, kernel_size_list, padding_list, and downsample_list must have the same length.") + + self.channel_list = [input_dim[0]] + channel_list + + # Initialize convolutional layers + layers = [] + for i in range(len(self.channel_list) - 1): + layers.append( + ConvSingleLayer( + in_channels=self.channel_list[i], + out_channels=self.channel_list[i + 1], + kernel_size=kernel_size_list[i], + padding=padding_list[i], + downsample=downsample_list[i] + ) + ) + + self.conv_layers = nn.ModuleList(layers) + + if penultimate_dim == 0: + # Try to auto compute penultimate_dim by processing the dummy input + with torch.no_grad(): + dummy_input = torch.randn(1, *input_dim) # B, C, H, W + x = dummy_input + for layer in self.conv_layers: + x = layer(x) + checked_penultimate_dim = x.numel() # Number of elements + else: + checked_penultimate_dim = penultimate_dim + + self.final_dense = nn.Linear(checked_penultimate_dim, output_dim, bias=True) + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + """Forward pass of the model. + + Args: + inputs (torch.Tensor): Input tensor of shape `(batch_size, C, H, W)`. + + Returns: + torch.Tensor: Output tensor of shape `(batch_size, output_dim)`. + + Raises: + ValueError: If inputs is not a rank 4 tensor. + """ + 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 + + for layer in self.conv_layers: + x = layer(x) + + x = torch.flatten(x, start_dim=1) # Flatten except batch-axis + x = self.final_dense(x) + return x diff --git a/modules/pytorch_wrapper/ptn_linear_model.py b/modules/pytorch_wrapper/ptn_linear_model.py index 3f1c9f5..bd2108e 100644 --- a/modules/pytorch_wrapper/ptn_linear_model.py +++ b/modules/pytorch_wrapper/ptn_linear_model.py @@ -10,7 +10,7 @@ class PtnLinearModel: A linear model consisting of dense layers. ### Fields: - - `dim_list`: A comma-separated string representing layer dimensions. + - `dim_list`: A string representing a Python list of layer dimensions. It should contain the input dimension for the first layer and the output dimensions for all layers. @@ -18,10 +18,11 @@ class PtnLinearModel: - Input data: 784 - Layer 1: 784 → 128 - Layer 2: 128 → 10 - - Specify `"784, 128, 10"` + - Specify `"[784,128,10]"` + + - `bias_list`: A string representing a Python list of bias usage (True/False) per layer. + Example: "[True,False]" - - `bias_list`: A comma-separated string representing bias usage (True/False) per layer. - - `num_layers`: The number of layers in the model. category: PyTorch wrapper - Training diff --git a/pyproject.toml b/pyproject.toml index 41c0000..50e24d1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "data-analysis" description = "Data analysis custom modules for ComfyUI – Integrates Pandas, Seaborn and Matplotlib into ComfyUI, featuring extensive PyTorch wrapper nodes for seamless tensor operations. Experimental PyTorch model training was added in v2.0." -version = "2.0.3" +version = "2.0.4" license = {file = "LICENSE"} [project.urls] diff --git a/tests/pytorch_wrapper/pt_evaluate_classification_conv_model_test.py b/tests/pytorch_wrapper/pt_evaluate_classification_conv_model_test.py new file mode 100644 index 0000000..b0d30cd --- /dev/null +++ b/tests/pytorch_wrapper/pt_evaluate_classification_conv_model_test.py @@ -0,0 +1,83 @@ +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_conv_model import PtnConvModel +from pytorch_wrapper.ptv_dataset import PtvDataset +from pytorch_wrapper.ptv_transforms_to_tensor import PtvTransformsToTensor +from pytorch_wrapper.pt_data_loader import PtDataLoader +from pytorch_wrapper.pt_evaluate_classification_model import PtEvaluateClassificationModel +from pytorch_wrapper.pt_load_model import PtLoadModel + +class TestPtEvaluate(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + # Transform + compose_node = PtvTransformsToTensor() + transform = compose_node.f()[0] + + # Dataset + dataset_node = PtvDataset() + dataset = dataset_node.f("FashionMNIST", + True, # download + "", # root + '{"train": False}', + transform=transform)[0] + + # Data loader + data_loader_node = PtDataLoader() + self.data_loader = data_loader_node.f( + dataset=dataset, + batch_size=32, + shuffle=True, + parameters='{"num_workers":1}' + )[0] + + # Model + self.model_node = PtnConvModel() + self.model = self.model_node.f( + input_dim="(1, 28, 28)", + penultimate_dim=0, + output_dim=10, + channel_list="[32,64,128,256]", + kernel_size_list="[3,3,3,3]", + padding_list="[1,1,1,1]", + downsample_list="[True,True,True,True]" + )[0] + + # Load weights + self.load_model = PtLoadModel() + epochs = 1 + self.loaded_model = self.load_model.f(self.model, f"fashion_mnist_{epochs}_epochs_conv.pt") + + # Instntiates evaluation node + self.evaluate = PtEvaluateClassificationModel() + + def test_evaluation(self): + """ + Tests evaluation. + + Note that this test requires a trained model file on the file system. + Unit test will fail if you run this for the first time. + Run the unit test for training to generate the weight, and re-run this unit test. + """ + accuracy, precision, recall, f1, num_samples = \ + self.evaluate.f(self.model, self.data_loader, use_gpu=True) + print(f""" + accuracy: {accuracy} + precision: {precision} + recall: {recall} + f1: {f1} + Number of samples: {num_samples} + """) + self.assertTrue(accuracy > 0.8) + self.assertEqual(num_samples, 10000) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/pt_train_classification_conv_model_test.py b/tests/pytorch_wrapper/pt_train_classification_conv_model_test.py new file mode 100644 index 0000000..295753a --- /dev/null +++ b/tests/pytorch_wrapper/pt_train_classification_conv_model_test.py @@ -0,0 +1,63 @@ +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_conv_model import PtnConvModel +from pytorch_wrapper.pto_adam import PtoAdam +from pytorch_wrapper.ptv_dataset import PtvDataset +from pytorch_wrapper.ptv_transforms_to_tensor import PtvTransformsToTensor +from pytorch_wrapper.pt_data_loader import PtDataLoader +from pytorch_wrapper.pt_train_classification_model import PtTrainClassificationModel +from pytorch_wrapper.pt_save_model import PtSaveModel + + +class TestPtTrainConvModel(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + compose_node = PtvTransformsToTensor() + transform = compose_node.f()[0] + dataset_node = PtvDataset() + dataset = dataset_node.f("FashionMNIST", + True, # download + "", # root + '{"train": True}', + transform=transform)[0] + data_loader_node = PtDataLoader() + self.train_loader = data_loader_node.f( + dataset=dataset, + batch_size=32, + shuffle=True, + parameters='{"num_workers":1}' + )[0] + + self.model_node = PtnConvModel() + self.model = self.model_node.f( + input_dim="(1, 28, 28)", + penultimate_dim=0, + output_dim=10, + channel_list="[32,64,128,256]", + kernel_size_list="[3,3,3,3]", + padding_list="[1,1,1,1]", + downsample_list="[True,True,True,True]" + )[0] + + self.optimizer = PtoAdam().f(self.model, 0.001, 0.9, 0.999)[0] + self.trainer = PtTrainClassificationModel() + self.save_model = PtSaveModel() + + def test_1(self): + epochs = 1 + trained_model = self.trainer.f(self.model, self.train_loader, self.optimizer, + epochs, # epochs + use_gpu=True)[0] + self.save_model.f(trained_model, f"fashion_mnist_{epochs}_epochs_conv.pt") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptn_conv_model_def_test.py b/tests/pytorch_wrapper/ptn_conv_model_def_test.py new file mode 100644 index 0000000..2b3b812 --- /dev/null +++ b/tests/pytorch_wrapper/ptn_conv_model_def_test.py @@ -0,0 +1,158 @@ +import os +import sys +import unittest +import torch + +PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", "..")) +MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules") +sys.path.append(MODULE_ROOT) +from pytorch_wrapper.ptn_conv_model_def import ConvModel, ConvSingleLayer + + +class TestPtnConvModelDef(unittest.TestCase): + def setUp(self): + """Setup method for initializing test conditions.""" + pass + + def test_single_layer_1_with_padding(self): + """Test ConvSingleLayer with padding=1, no downsampling.""" + model = ConvSingleLayer( + in_channels=4, + out_channels=32, + kernel_size=3, + padding=1, + downsample=False, + apply_activation=True) + x = torch.ones(1, 4, 28, 28) + out = model(x) + self.assertEqual(out.size(), torch.Size([1, 32, 28, 28])) + + def test_single_layer_no_padding(self): + """Test ConvSingleLayer with padding=0, no downsampling.""" + model = ConvSingleLayer( + in_channels=4, + out_channels=32, + kernel_size=3, + padding=0, + downsample=False, + apply_activation=True) + x = torch.ones(1, 4, 28, 28) + out = model(x) + self.assertEqual(out.size(), torch.Size([1, 32, 26, 26])) + + def test_single_layer_3_no_padding_downsample(self): + """Test ConvSingleLayer with padding=0 and downsampling enabled.""" + model = ConvSingleLayer( + in_channels=4, + out_channels=32, + kernel_size=3, + padding=0, + downsample=True, + apply_activation=True) + x = torch.ones(1, 4, 28, 28) + out = model(x) + self.assertEqual(out.size(), torch.Size([1, 32, 13, 13])) + + def test_1(self): + """Test ConvModel with a single convolutional layer, checking output size.""" + model = ConvModel( + input_dim=(4, 28, 28), + penultimate_dim=0, + output_dim=10, + channel_list=[32], + kernel_size_list=[3], + padding_list=[1], + downsample_list=[False]) + + x = torch.ones(1, 4, 28, 28) + out = model(x) + self.assertEqual(out.size(), torch.Size([1, 10])) + + def test_2(self): + """Test ConvModel intermediate convolution layers' output size.""" + model = ConvModel( + input_dim=(4, 28, 28), + penultimate_dim=0, + output_dim=10, + channel_list=[32], + kernel_size_list=[3], + padding_list=[1], + downsample_list=[False]) + + x = torch.ones(1, 4, 28, 28) + for c in model.conv_layers: + x = c(x) + out = x + self.assertEqual(out.size(), torch.Size([1, 32, 28, 28])) + + def test_3(self): + """Test ConvModel with two convolutional layers and downsampling enabled.""" + model = ConvModel( + input_dim=(4, 28, 28), + penultimate_dim=0, + output_dim=10, + channel_list=[32,64], + kernel_size_list=[3,3], + padding_list=[1,1], + downsample_list=[True,True]) + + x = torch.ones(1, 4, 28, 28) + for c in model.conv_layers: + x = c(x) + out = x + self.assertEqual(out.size(), torch.Size([1, 64, 7, 7])) + + def test_4(self): + """Test ConvModel with three convolutional layers and downsampling enabled.""" + model = ConvModel( + input_dim=(4, 28, 28), + penultimate_dim=0, + output_dim=10, + channel_list=[32,64,128], + kernel_size_list=[3,3,3], + padding_list=[1,1,1], + downsample_list=[True,True,True]) + + x = torch.ones(1, 4, 28, 28) + for c in model.conv_layers: + x = c(x) + out = x + self.assertEqual(out.size(), torch.Size([1, 128, 3, 3])) + + def test_5(self): + """Test ConvModel with four convolutional layers and downsampling enabled.""" + model = ConvModel( + input_dim=(4, 28, 28), + penultimate_dim=0, + output_dim=10, + channel_list=[32,64,128,256], + kernel_size_list=[3,3,3,3], + padding_list=[1,1,1,1], + downsample_list=[True,True,True,True]) + + x = torch.ones(1, 4, 28, 28) + for c in model.conv_layers: + x = c(x) + out = x + self.assertEqual(out.size(), torch.Size([1, 256, 1, 1])) + + def test_6(self): + """Test ConvModel with five convolutional layers, last layer without downsampling.""" + model = ConvModel( + input_dim=(4, 28, 28), + penultimate_dim=0, + output_dim=10, + channel_list=[32,64,128,256,512], + kernel_size_list=[3,3,3,3,1], + padding_list=[1,1,1,1,0], + downsample_list=[True,True,True,True,False]) + + x = torch.ones(1, 4, 28, 28) + for c in model.conv_layers: + x = c(x) + out = x + self.assertEqual(out.size(), torch.Size([1, 512, 1, 1])) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptn_conv_model_test.py b/tests/pytorch_wrapper/ptn_conv_model_test.py new file mode 100644 index 0000000..9e6564e --- /dev/null +++ b/tests/pytorch_wrapper/ptn_conv_model_test.py @@ -0,0 +1,64 @@ +import os +import sys +import unittest +import torch + +PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", "..")) +MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules") +sys.path.append(MODULE_ROOT) +from pytorch_wrapper.ptn_conv_model import PtnConvModel + + +class TestPtnConvModel(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.model_node = PtnConvModel() + + def test_forward(self): + model = self.model_node.f( + input_dim="(4, 28, 28)", + penultimate_dim=0, + output_dim=10, + channel_list="[32,64,128,256]", + kernel_size_list="[3,3,3,3]", + padding_list="[1,1,1,1]", + downsample_list="[True,True,True,True]" + )[0] + bs = 2 + x = torch.ones(bs, 4, 28, 28) + out = model(x) + self.assertEqual(out.size(), torch.Size([bs, 10])) + + def test_rank3_input(self): + model = self.model_node.f( + input_dim="(1, 28, 28)", + penultimate_dim=0, + output_dim=10, + channel_list="[32,64,128,256]", + kernel_size_list="[3,3,3,3]", + padding_list="[1,1,1,1]", + downsample_list="[True,True,True,True]" + )[0] + bs = 2 + x = torch.ones(bs, 28, 28) + out = model(x) + self.assertEqual(out.size(), torch.Size([bs, 10])) + + def test_rank2_input(self): + model = self.model_node.f( + input_dim="(1, 28, 28)", + penultimate_dim=0, + output_dim=10, + channel_list="[32,64,128,256]", + kernel_size_list="[3,3,3,3]", + padding_list="[1,1,1,1]", + downsample_list="[True,True,True,True]" + )[0] + bs = 2 + x = torch.ones(bs, 28) + with self.assertRaises(ValueError) as context: + model(x) + + +if __name__ == "__main__": + unittest.main()