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

Added conv model for pytorch training.

Browse Source
This commit is contained in:
HowToSD
2025-03-04 23:53:50 -08:00
parent 4a23816149
commit b881244d0c
21 changed files with 725 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+26 -8
View File
@@ -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.
<hr>
__init__.py
+3
View File
@@ -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",
docs/images/conv_eval.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 177 KiB

docs/images/conv_train.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 194 KiB

docs/images/conv_train2.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 69 KiB

docs/reference/node_reference.md
+2 -4
View File
@@ -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. |
docs/reference/pt_flatten.md
+1 -3
View File
@@ -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 |
<HR>
Category: Miscellaneous
Category: PyTorch wrapper - Transform
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptn_conv_model.md
+45
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Training
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptn_conv_model_def.md
+17
View File
@@ -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.
<HR>
Category: Miscellaneous
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptn_linear_model.md
+4 -3
View File
@@ -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.
examples/workflows/fashion_mnist_eval_conv.json
+1
View File
@@ -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}
examples/workflows/fashion_mnist_train_conv.json
+1
View File
@@ -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}
modules/pytorch_wrapper/pt_flatten.py
+1 -1
View File
@@ -6,7 +6,7 @@ class PtFlatten:
"""
Flattens a PyTorch tensor into a 1D tensor.
Category: PyTorch wrapper
category: PyTorch wrapper - Transform
"""
@classmethod
modules/pytorch_wrapper/ptn_conv_model.py
+99
View File
@@ -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,)
modules/pytorch_wrapper/ptn_conv_model_def.py
+151
View File
@@ -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
modules/pytorch_wrapper/ptn_linear_model.py
+5 -4
View File
@@ -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
pyproject.toml
+1 -1
View File
@@ -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]
tests/pytorch_wrapper/pt_evaluate_classification_conv_model_test.py
+83
View File
@@ -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()
tests/pytorch_wrapper/pt_train_classification_conv_model_test.py
+63
View File
@@ -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()
tests/pytorch_wrapper/ptn_conv_model_def_test.py
+158
View File
@@ -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()
tests/pytorch_wrapper/ptn_conv_model_test.py
+64
View File
@@ -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()