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

Added experimental PyTorch model training.

Browse Source
This commit is contained in:
HowToSD
2025-03-03 18:39:52 -08:00
parent 62e4f6efb8
commit 11d3c9490c
40 changed files with 1649 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+3 -1
View File
@@ -1,3 +1,5 @@
__pycache__/
**/__pycache__/
datasets/
models/
*.pth
README.md
+24 -1
View File
@@ -1,4 +1,27 @@
# ComfyUI-Data-Analysis
## Update (March 3, 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:
Training: examples/fashion_mnist_train.json
Evaluation: examples/fashion_mnist_eval.json
This workflow trains a FashionMNIST model from scratch effortlessly.
**Training workflow**
![Train](docs/images/fashion_mnist_train.png)
**Eval workflow**
![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.
<hr>
![Baseball tutorial](docs/images/baseball_question.jpg)
## How Can We Figure Out Which MLB Player Had the Most Hits Per Year?
@@ -61,7 +84,7 @@ Note: GPU is not required to run this custom module.
This will create a new subdirectory ComfyUI-Data-Analysis.
3. Rename the folder name from ComfyUI-Data-Analysis to data-analysis.
Skipping this step will cause example workflows to fail, as they rely on this folder name.
4. Check if your ComfyUI environment already has pandas, matplotlib, scipy and seaborn. If not, install them using pip.
4. Check if your ComfyUI environment already has required packages listed in requirements.txt. If not, install them using pip.
```
pip install -r requirements.txt
```
__init__.py
+30
View File
@@ -141,9 +141,11 @@ from .modules.pytorch_wrapper.pt_bmm import PtBmm
from .modules.pytorch_wrapper.pt_bool_create import PtBoolCreate
from .modules.pytorch_wrapper.pt_cos import PtCos
from .modules.pytorch_wrapper.pt_cosh import PtCosh
from .modules.pytorch_wrapper.pt_data_loader import PtDataLoader
from .modules.pytorch_wrapper.pt_div import PtDiv
from .modules.pytorch_wrapper.pt_einsum import PtEinsum
from .modules.pytorch_wrapper.pt_eq import PtEq
from .modules.pytorch_wrapper.pt_evaluate_classification_model import PtEvaluateClassificationModel
from .modules.pytorch_wrapper.pt_exp import PtExp
from .modules.pytorch_wrapper.pt_flatten import PtFlatten
from .modules.pytorch_wrapper.pt_float_create import PtFloatCreate
@@ -159,6 +161,7 @@ from .modules.pytorch_wrapper.pt_index_select import PtIndexSelect
from .modules.pytorch_wrapper.pt_int_create import PtIntCreate
from .modules.pytorch_wrapper.pt_le import PtLe
from .modules.pytorch_wrapper.pt_linspace import PtLinspace
from .modules.pytorch_wrapper.pt_load_model import PtLoadModel
from .modules.pytorch_wrapper.pt_log import PtLog
from .modules.pytorch_wrapper.pt_logical_and import PtLogicalAnd
from .modules.pytorch_wrapper.pt_logical_not import PtLogicalNot
@@ -184,6 +187,7 @@ from .modules.pytorch_wrapper.pt_rand_int import PtRandInt
from .modules.pytorch_wrapper.pt_randn import PtRandn
from .modules.pytorch_wrapper.pt_remainder import PtRemainder
from .modules.pytorch_wrapper.pt_reshape import PtReshape
from .modules.pytorch_wrapper.pt_save_model import PtSaveModel
from .modules.pytorch_wrapper.pt_scatter import PtScatter
from .modules.pytorch_wrapper.pt_show_size import PtShowSize
from .modules.pytorch_wrapper.pt_show_text import PtShowText
@@ -204,11 +208,17 @@ from .modules.pytorch_wrapper.pt_to_image import PtToImage
from .modules.pytorch_wrapper.pt_to_latent import PtToLatent
from .modules.pytorch_wrapper.pt_to_numpy import PtToNumpy
from .modules.pytorch_wrapper.pt_to_rgb_tensors import PtToRgbTensors
from .modules.pytorch_wrapper.pt_train_classification_model import PtTrainClassificationModel
from .modules.pytorch_wrapper.pt_unsqueeze import PtUnsqueeze
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_linear_model import PtnLinearModel
from .modules.pytorch_wrapper.pto_adam import PtoAdam
from .modules.pytorch_wrapper.ptv_dataset import PtvDataset
from .modules.pytorch_wrapper.ptv_dataset_loader import PtvDatasetLoader
from .modules.pytorch_wrapper.ptv_transforms_to_tensor import PtvTransformsToTensor
from .modules.seaborn_wrapper.sns_bar import SNSBar
from .modules.seaborn_wrapper.sns_boxplot import SNSBoxplot
from .modules.seaborn_wrapper.sns_heatmap import SNSHeatmap
@@ -363,9 +373,11 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
"PtBoolCreate": PtBoolCreate,
"PtCos": PtCos,
"PtCosh": PtCosh,
"PtDataLoader": PtDataLoader,
"PtDiv": PtDiv,
"PtEinsum": PtEinsum,
"PtEq": PtEq,
"PtEvaluateClassificationModel": PtEvaluateClassificationModel,
"PtExp": PtExp,
"PtFlatten": PtFlatten,
"PtFloatCreate": PtFloatCreate,
@@ -381,6 +393,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
"PtIntCreate": PtIntCreate,
"PtLe": PtLe,
"PtLinspace": PtLinspace,
"PtLoadModel": PtLoadModel,
"PtLog": PtLog,
"PtLogicalAnd": PtLogicalAnd,
"PtLogicalNot": PtLogicalNot,
@@ -406,6 +419,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
"PtRandn": PtRandn,
"PtRemainder": PtRemainder,
"PtReshape": PtReshape,
"PtSaveModel": PtSaveModel,
"PtScatter": PtScatter,
"PtShowSize": PtShowSize,
"PtShowText": PtShowText,
@@ -426,11 +440,17 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
"PtToLatent": PtToLatent,
"PtToNumpy": PtToNumpy,
"PtToRgbTensors": PtToRgbTensors,
"PtTrainClassificationModel": PtTrainClassificationModel,
"PtUnsqueeze": PtUnsqueeze,
"PtVar": PtVar,
"PtView": PtView,
"PtWhere": PtWhere,
"PtZeros": PtZeros,
"PtnLinearModel": PtnLinearModel,
"PtoAdam": PtoAdam,
"PtvDataset": PtvDataset,
"PtvDatasetLoader": PtvDatasetLoader,
"PtvTransformsToTensor": PtvTransformsToTensor,
"PyKvFloatCreate": PyKvFloatCreate,
"PyKvIntCreate": PyKvIntCreate,
"PyKvStringCreate": PyKvStringCreate,
@@ -588,9 +608,11 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
"PtBoolCreate": "Pt Bool Create",
"PtCos": "Pt Cos",
"PtCosh": "Pt Cosh",
"PtDataLoader": "Pt Data Loader",
"PtDiv": "Pt Div",
"PtEinsum": "Pt Einsum",
"PtEq": "Pt Eq",
"PtEvaluateClassificationModel": "Pt Evaluate Classification Model",
"PtExp": "Pt Exp",
"PtFlatten": "Pt Flatten",
"PtFloatCreate": "Pt Float Create",
@@ -606,6 +628,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
"PtIntCreate": "Pt Int Create",
"PtLe": "Pt Le",
"PtLinspace": "Pt Linspace",
"PtLoadModel": "Pt Load Model",
"PtLog": "Pt Log",
"PtLogicalAnd": "Pt Logical And",
"PtLogicalNot": "Pt Logical Not",
@@ -631,6 +654,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
"PtRandn": "Pt Randn",
"PtRemainder": "Pt Remainder",
"PtReshape": "Pt Reshape",
"PtSaveModel": "Pt Save Model",
"PtScatter": "Pt Scatter",
"PtShowSize": "Pt Show Size",
"PtShowText": "Pt Show Text",
@@ -651,11 +675,17 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
"PtToLatent": "Pt To Latent",
"PtToNumpy": "Pt To Numpy",
"PtToRgbTensors": "Pt To Rgb Tensors",
"PtTrainClassificationModel": "Pt Train Classification Model",
"PtUnsqueeze": "Pt Unsqueeze",
"PtVar": "Pt Var",
"PtView": "Pt View",
"PtWhere": "Pt Where",
"PtZeros": "Pt Zeros",
"PtnLinearModel": "Ptn Linear Model",
"PtoAdam": "Pto Adam",
"PtvDataset": "Ptv Dataset",
"PtvDatasetLoader": "Ptv Dataset Loader",
"PtvTransformsToTensor": "Ptv Transforms To Tensor",
"PyKvFloatCreate": "Py Kv Float Create",
"PyKvIntCreate": "Py Kv Int Create",
"PyKvStringCreate": "Py Kv String Create",
docs/images/fashion_mnist_eval.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 119 KiB

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

After

Width:  |  Height:  |  Size: 165 KiB

docs/reference/node_reference.md
+13
View File
@@ -113,6 +113,7 @@
| [Pandas Load JSON](pandas_load_json.md) | Loads JSON files into a pandas DataFrame. |
| [Pandas Save CSV](pandas_save_csv.md) | Saves a PandasDataFrame to a CSV file. |
| [Pandas Save JSON](pandas_save_json.md) | Saves a PandasDataFrame to a JSON file. |
| [Pt Save Model](pt_save_model.md) | A wrapper class for saving a PyTorch model. |
## Math
| Node | Description |
| --- | --- |
@@ -257,6 +258,18 @@
| [Pt To Latent](pt_to_latent.md) | Casts a PyTorch tensor as a latent tensor. |
| [Pt To Numpy](pt_to_numpy.md) | Converts PyTorch tensor to NumPy ndarray. |
| [Pt To Rgb Tensors](pt_to_rgb_tensors.md) | Splits a PyTorch tensor into R, G, and B tensors. |
| [Pto Adam](pto_adam.md) | Converts PyTorch tensor to NumPy ndarray. |
## PyTorch wrapper - Training
| Node | Description |
| --- | --- |
| [Pt Data Loader](pt_data_loader.md) | Loads data from a dataset node and creates a PyTorch DataLoader. |
| [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 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. |
| [Ptv Transforms To Tensor](ptv_transforms_to_tensor.md) | Transforms elements of dataset to PyTorch tensors. |
## PyTorch wrapper - Transform
| Node | Description |
| --- | --- |
docs/reference/pt_data_loader.md
+29
View File
@@ -0,0 +1,29 @@
# Pt Data Loader
Loads data from a dataset node and creates a PyTorch DataLoader.
**dataset**
Connect a PtvDataset node that provides the dataset.
**batch size**
Specifies the number of samples per batch.
**shuffle**
If True, shuffles the dataset before each epoch.
## Input
| Name | Data type |
|---|---|
| dataset | Ptvdataset |
| batch_size | Int |
| shuffle | Boolean |
| parameters | String |
## Output
| Data type |
|---|
| Ptdataloader |
<HR>
Category: PyTorch wrapper - Training
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pt_evaluate_classification_model.md
+23
View File
@@ -0,0 +1,23 @@
# Pt Evaluate Classification Model
Performs inference on test data and computes evaluation metrics.
## Input
| Name | Data type |
|---|---|
| model | Ptmodel |
| data_loader | Ptdataloader |
| use_gpu | Boolean |
## Output
| Data type |
|---|
| Float |
| Float |
| Float |
| Float |
| Int |
<HR>
Category: PyTorch wrapper - Training
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pt_load_model.md
+18
View File
@@ -0,0 +1,18 @@
# Pt Load Model
A wrapper class for saving a PyTorch model.
## Input
| Name | Data type |
|---|---|
| model | Ptmodel |
| model_path | 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/pt_save_model.md
+18
View File
@@ -0,0 +1,18 @@
# Pt Save Model
A wrapper class for saving a PyTorch model.
## Input
| Name | Data type |
|---|---|
| model | Ptmodel |
| model_path | String |
## Output
| Data type |
|---|
| String |
<HR>
Category: IO
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pt_train_classification_model.md
+24
View File
@@ -0,0 +1,24 @@
# Pt Train Classification Model
Trains a classification model using a given dataset, optimizer, and number of epochs.
## Input
| Name | Data type |
|---|---|
| model | Ptmodel |
| train_loader | Ptdataloader |
| optimizer | Ptoptimizer |
| epochs | Int |
| use_gpu | Boolean |
| val_loader | Ptdataloader |
## Output
| Data type |
|---|
| Ptmodel |
| Tensor |
| Tensor |
<HR>
Category: PyTorch wrapper - Training
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptn_linear_model.md
+34
View File
@@ -0,0 +1,34 @@
# Ptn Linear Model
A linear model consisting of dense layers.
### Fields:
- `dim_list`: A comma-separated string representing layer dimensions.
It should contain the input dimension for the first layer and the output dimensions
for all layers.
Example for a two-layer network:
- Input data: 784
- Layer 1: 784 → 128
- Layer 2: 128 → 10
- Specify `"784, 128, 10"`
- `bias_list`: A comma-separated string representing bias usage (True/False) per layer.
- `num_layers`: The number of layers in the model.
## Input
| Name | Data type |
|---|---|
| dim_list | String |
| bias_list | String |
| num_layers | Int |
## 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_linear_model_def.md
+20
View File
@@ -0,0 +1,20 @@
# Dense Model
A multi-layer dense neural network with ReLU activation after each dense layer, except for the final layer.
Args:
dim_list (List[int]): List of feature dimensions. Should be of length `num_layers + 1`.
bias_list (List[bool]): List indicating whether each layer should have a bias.
num_layers (int): Number of layers in the model.
Example:
For a three-layer network:
- Layer 1: input -> output1
- Layer 2: output1 -> output2
- Layer 3: output2 -> output3
`dim_list = [input_dim, output1_dim, output2_dim, output3_dim]`
<HR>
Category: Miscellaneous
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pto_adam.md
+20
View File
@@ -0,0 +1,20 @@
# Pto Adam
Converts PyTorch tensor to NumPy ndarray.
## Input
| Name | Data type |
|---|---|
| model | Ptmodel |
| learning_rate | Float |
| beta1 | Float |
| beta2 | Float |
## Output
| Data type |
|---|
| Ptoptimizer |
<HR>
Category: PyTorch wrapper - Tensor data conversion
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptv_dataset.md
+51
View File
@@ -0,0 +1,51 @@
# Ptv Dataset
A Torchvision Dataset class wrapper.
Parameters
**name**
Specify a dataset class name such as MNIST, FashionMNIST in string.
This will be converted to a dataset class internally.
**download**
Set to True if you want to download.
**root**
Specify the root directory of the downloaded dataset relative to the
dataset directory of this extension, or specify the absolute path.
If you leave as a blank, the dataset will be downloaded under the `datasets` directory
of this extension.
**transform**
Data transforms e.g. transform to PyTorch tensor, normalize image.
Plug in the PTV Transforms node that contains Torchvision Transforms functionality.
**parameters**
Specify other parameters in Python dict format.
For example, if you want to set Train parameter to True to download the train set, specify:
```
{"train": True}
```
or if you want to download the test set, specify:
```
{"train": False}
```
## Input
| Name | Data type |
|---|---|
| name | String |
| download | Boolean |
| root | String |
| transform | Ptvtransform |
| parameters | String |
## Output
| Data type |
|---|
| Ptvdataset |
<HR>
Category: PyTorch wrapper - Training
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptv_dataset_loader.md
+60
View File
@@ -0,0 +1,60 @@
# Ptv Dataset Loader
A node to combine the dataset and data loader into a single node.
Transform is internally called to convert input data to PyTorch tensors.
Parameters
**name**
Specify the dataset class name such as MNIST, FashionMNIST in string.
This will be converted to a dataset class internally.
**download**
Set to True if you want to download.
**root**
Specify the root directory of the downloaded dataset relative to the
dataset directory of this extension, or specify the absolute path.
If you leave as a blank, the dataset will be downloaded under the `datasets` directory
of this extension.
**batch_size**
The number of samples per batch.
**shuffle**
If True, shuffles the dataset before each epoch.
**dataset parameters**
Specify other parameters in Python dict format.
For example, if you want to set Train parameter to True to download the train set, specify:
```
{"train": True}
```
or if you want to download the test set, specify:
```
{"train": False}
```
**load parameters**
Additional parameters in Python dictionary format.
## Input
| Name | Data type |
|---|---|
| name | String |
| download | Boolean |
| root | String |
| batch_size | Int |
| shuffle | Boolean |
| dataset_parameters | String |
| load_parameters | String |
## Output
| Data type |
|---|
| Ptdataloader |
<HR>
Category: PyTorch wrapper - Training
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptv_transforms_to_tensor.md
+12
View File
@@ -0,0 +1,12 @@
# Ptv Transforms To Tensor
Transforms elements of dataset to PyTorch tensors.
## Output
| Data type |
|---|
| Ptvtransform |
<HR>
Category: PyTorch wrapper - Training
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
examples/workflows/fashion_mnist_eval.json
+1
View File
@@ -0,0 +1 @@
{"last_node_id":47,"last_link_id":61,"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":22,"type":"PtnLinearModel","pos":[2886.657958984375,1718.1610107421875],"size":[264.3730163574219,109.75016021728516],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[24],"slot_index":0}],"properties":{"Node name for S&R":"PtnLinearModel"},"widgets_values":["[784,180,42,10]","[False,False,False]",3]},{"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.8705378412983575"]},{"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.8721685257742706"]},{"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.8715999999999999"]},{"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.8716"]},{"id":46,"type":"PtvDatasetLoader","pos":[2884.05078125,1908.8636474609375],"size":[266.8697814941406,246],"flags":{},"order":1,"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":23,"type":"PtLoadModel","pos":[3243.927978515625,1718.1610107421875],"size":[397.5032653808594,67.37537384033203],"flags":{},"order":2,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":24}],"outputs":[{"name":"model","type":"PTMODEL","links":[55],"slot_index":0}],"properties":{"Node name for S&R":"PtLoadModel"},"widgets_values":["fashion_mnist_ep5.pt"]}],"links":[[24,22,0,23,0,"PTMODEL"],[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"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.8888542618071794,"offset":[-2627.7528215343027,-1537.849355050907]},"node_versions":{"data-analysis":"dab52eaa403a5e9cdc546aca7f7d95e41ea0a943"},"ue_links":[]},"version":0.4}
examples/workflows/fashion_mnist_train.json
+1
View File
@@ -0,0 +1 @@
{"last_node_id":56,"last_link_id":71,"nodes":[{"id":22,"type":"PtnLinearModel","pos":[2830.718017578125,1445.9212646484375],"size":[222.090087890625,128.31336975097656],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[22,65],"slot_index":0}],"properties":{"Node name for S&R":"PtnLinearModel"},"widgets_values":["[784,180,42,10]","[False,False,False]",3]},{"id":47,"type":"PtvDatasetLoader","pos":[2830.718017578125,1652.3197021484375],"size":[239.68125915527344,246],"flags":{},"order":1,"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":48,"type":"PtvDatasetLoader","pos":[2830.718017578125,1957.211181640625],"size":[239.68125915527344,246],"flags":{},"order":2,"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":50,"type":"PtTrainClassificationModel","pos":[3170.615478515625,1736.97705078125],"size":[243.70773315429688,142],"flags":{},"order":9,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":65},{"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":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_ep5.pt"]},{"id":54,"type":"MarkdownNote","pos":[2582.973388671875,1955.99560546875],"size":[210,58.94380569458008],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Load val data"],"color":"#432","bgcolor":"#653"},{"id":53,"type":"MarkdownNote","pos":[2582.973388671875,1652.8607177734375],"size":[210,58.94380569458008],"flags":{},"order":4,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Load train data"],"color":"#432","bgcolor":"#653"},{"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":22}],"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":52,"type":"MarkdownNote","pos":[2582.973388671875,1445.9212646484375],"size":[210,58.94380569458008],"flags":{},"order":5,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Define model"],"color":"#432","bgcolor":"#653"},{"id":55,"type":"MarkdownNote","pos":[3170.615478515625,1953.1151123046875],"size":[210,58.94380569458008],"flags":{},"order":6,"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":7,"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_ep5.pt"]}],"links":[[22,22,0,11,0,"PTMODEL"],[65,22,0,50,0,"PTMODEL"],[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"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.8678754334266201,"offset":[-2457.1876923760833,-1312.0857787454647]},"node_versions":{"data-analysis":"dab52eaa403a5e9cdc546aca7f7d95e41ea0a943"},"ue_links":[]},"version":0.4}
modules/pytorch_wrapper/pt_data_loader.py
+81
View File
@@ -0,0 +1,81 @@
import inspect
from typing import Any, Dict
import ast
import torch
from torch.utils.data import DataLoader
class PtDataLoader:
"""
Loads data from a dataset node and creates a PyTorch DataLoader.
**dataset**
Connect a PtvDataset node that provides the dataset.
**batch size**
Specifies the number of samples per batch.
**shuffle**
If True, shuffles the dataset before each epoch.
category: PyTorch wrapper - Training
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types required for data loading.
Returns:
Dict[str, Any]: A dictionary specifying required and optional input types.
"""
return {
"required": {
"dataset": ("PTVDATASET", {}),
"batch_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
"shuffle": ("BOOLEAN", {"default": False}),
},
"optional": {
"parameters": ("STRING", {"default": "", "multiline": True})
}
}
RETURN_TYPES: tuple = ("PTDATALOADER",)
FUNCTION: str = "f"
CATEGORY: str = "Training"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self,
dataset: Any, # One of the classes defined in torchvision.datasets
batch_size: int,
shuffle: bool,
parameters: str) -> tuple:
"""
Creates a PyTorch DataLoader for the given dataset.
Args:
dataset (torchvision.datasets.{MNIST, ...}): An instance of a dataset from torchvision.datasets.
batch_size (int): The number of samples per batch.
shuffle (bool): If True, shuffles the dataset before each epoch.
parameters (str): Additional parameters in Python dictionary format.
Returns:
tuple: A tuple containing the created DataLoader instance.
"""
with torch.inference_mode(False):
param_dict = ast.literal_eval(parameters)
# Get valid parameters for the DataLoader class
valid_params = inspect.signature(DataLoader).parameters
# Filter only parameters that match the dataset's constructor
filtered_params = {k: v for k, v in param_dict.items() if k in valid_params}
dl = DataLoader(dataset=dataset,
batch_size=batch_size,
shuffle=shuffle,
**filtered_params)
return (dl,)
modules/pytorch_wrapper/pt_evaluate_classification_model.py
+71
View File
@@ -0,0 +1,71 @@
from typing import Any, Dict, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
class PtEvaluateClassificationModel:
"""
Performs inference on test data and computes evaluation metrics.
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": {
"model": ("PTMODEL", {}),
"data_loader": ("PTDATALOADER", {}),
"use_gpu": ("BOOLEAN", {"default":False})
}
}
RETURN_NAMES: Tuple[str, ...] = ("accuracy", "precision", "recall", "f1", "num_samples")
RETURN_TYPES: Tuple[str, ...] = ("FLOAT", "FLOAT", "FLOAT", "FLOAT", "INT")
FUNCTION: str = "f"
CATEGORY: str = "Training"
def f(self, model: nn.Module, data_loader: torch.utils.data.DataLoader, use_gpu: bool) -> Tuple[float, float, float, float, int]:
"""
Performs inference on test data and computes evaluation metrics.
Args:
model (torch.nn.Module): The PyTorch model to evaluate.
data_loader (torch.utils.data.DataLoader): DataLoader for the test dataset.
use_gpu (bool): Whether to use GPU for inference.
Returns:
Tuple[float, float, float, float, int]: A tuple containing accuracy, precision, recall,
F1-score, and the number of samples.
"""
with torch.inference_mode():
if use_gpu:
model.to("cuda")
model.eval()
y_hat = []
y = []
for images, labels in data_loader:
if use_gpu:
images = images.to("cuda")
labels = labels.to("cuda")
outputs = model(images)
preds = outputs.argmax(dim=1)
y_hat.extend(preds.cpu().numpy())
y.extend(labels.cpu().numpy())
accuracy = accuracy_score(y, y_hat)
precision = precision_score(y, y_hat, average='macro')
recall = recall_score(y, y_hat, average='macro')
f1 = f1_score(y, y_hat, average='macro')
num_samples = len(y)
return (accuracy, precision, recall, f1, num_samples)
modules/pytorch_wrapper/pt_load_model.py
+41
View File
@@ -0,0 +1,41 @@
import os
from typing import Any, Dict
import torch
from .utils import get_model_full_path
class PtLoadModel:
"""
A wrapper class for saving a PyTorch model.
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": {
"model": ("PTMODEL", {}),
"model_path": ("STRING", {"default": "model.pkl"})
}
}
RETURN_NAMES: tuple = ("model",)
RETURN_TYPES: tuple = ("PTMODEL",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self, model:Any, model_path: str):
model_full_path = get_model_full_path(model_path)
sd = torch.load(model_full_path)
model.load_state_dict(sd)
return (model,)
modules/pytorch_wrapper/pt_save_model.py
+36
View File
@@ -0,0 +1,36 @@
import os
from typing import Any, Dict
import torch
from .utils import get_model_full_path
class PtSaveModel:
"""
A wrapper class for saving a PyTorch model.
category: IO
"""
@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": {
"model": ("PTMODEL", {}),
"model_path": ("STRING", {"default": "model.pkl"})
}
}
RETURN_NAMES: tuple = ("model full path",)
RETURN_TYPES: tuple = ("STRING",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, model:Any, model_path: str):
model_full_path = get_model_full_path(model_path)
torch.save(model.state_dict(), model_full_path)
return (os.path.realpath(model_full_path),)
modules/pytorch_wrapper/pt_train_classification_model.py
+111
View File
@@ -0,0 +1,111 @@
from typing import Any, Dict, Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
class PtTrainClassificationModel:
"""
Trains a classification model using a given dataset, optimizer, and number of epochs.
category: PyTorch wrapper - Training
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types required for training the model.
Returns:
Dict[str, Any]: A dictionary specifying required and optional input types.
"""
return {
"required": {
"model": ("PTMODEL", {}),
"train_loader": ("PTDATALOADER", {}),
"optimizer": ("PTOPTIMIZER", {}),
"epochs": ("INT", {"default":1, "min":1, "max":1e6}),
"use_gpu": ("BOOLEAN", {"default":False})
},
"optional": {
"val_loader": ("PTDATALOADER", {}),
}
}
RETURN_NAMES: tuple = ("Model", "train loss", "val loss")
RETURN_TYPES: tuple = ("PTMODEL","TENSOR","TENSOR")
FUNCTION: str = "f"
CATEGORY: str = "Training"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self, model: nn.Module,
train_loader: torch.utils.data.DataLoader,
optimizer: torch.optim.Optimizer,
epochs: int,
use_gpu: bool,
val_loader: Optional[torch.utils.data.DataLoader] = None) -> Tuple[nn.Module, torch.Tensor, torch.Tensor]:
"""
Trains a classification model.
Args:
model (torch.nn.Module): The PyTorch model to be trained.
train_loader (torch.utils.data.DataLoader): DataLoader containing the training dataset.
optimizer (torch.optim.Optimizer): Optimizer used for training.
epochs (int): Number of training epochs.
use_gpu (bool): Whether to use GPU for training.
val_loader (Optional[torch.utils.data.DataLoader]): DataLoader for validation dataset (optional).
Returns:
Tuple[torch.nn.Module, torch.Tensor, torch.Tensor]: A tuple containing the trained model,
training loss history, and validation loss history.
"""
with torch.inference_mode(False):
train_loss = list()
val_loss = list()
loss_function = nn.CrossEntropyLoss(reduction="sum")
# Train the model
for epoch in range(epochs):
if use_gpu:
model.to("cuda")
model.train()
total_loss = 0
for images, labels in train_loader:
if use_gpu:
images = images.to("cuda")
labels = labels.to("cuda")
optimizer.zero_grad()
outputs = model(images)
loss = loss_function(outputs, labels)
loss.backward()
optimizer.step()
total_loss += loss.item()
epoch_loss = total_loss / len(train_loader)
print(f"Epoch (train) {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}")
train_loss.append(epoch_loss)
if val_loader:
pass
with torch.no_grad():
model.eval()
total_loss = 0
for images, labels in val_loader:
if use_gpu:
images = images.to("cuda")
labels = labels.to("cuda")
outputs = model(images)
loss = loss_function(outputs, labels)
total_loss += loss.item()
epoch_loss = total_loss / len(val_loader)
print(f"Epoch (val) {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}")
val_loss.append(epoch_loss)
return (model,
torch.Tensor(np.asarray(train_loss)),
torch.Tensor(np.asarray(val_loss)))
modules/pytorch_wrapper/ptn_linear_model.py
+86
View File
@@ -0,0 +1,86 @@
from typing import Any, Dict, Tuple
import ast
import torch
import torch.nn as nn
from .ptn_linear_model_def import DenseModel
class PtnLinearModel:
"""
A linear model consisting of dense layers.
### Fields:
- `dim_list`: A comma-separated string representing layer dimensions.
It should contain the input dimension for the first layer and the output dimensions
for all layers.
Example for a two-layer network:
- Input data: 784
- Layer 1: 784 → 128
- Layer 2: 128 → 10
- Specify `"784, 128, 10"`
- `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
"""
@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": {
"dim_list": ("STRING", {"default": "784,10"}),
"bias_list": ("STRING", {"default": "{True}"}),
"num_layers": ("INT", {"default": 1, "min": 1, "max": 2000}),
}
}
RETURN_TYPES: Tuple[str] = ("PTMODEL",)
FUNCTION: str = "f"
CATEGORY: str = "Training"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self, dim_list: str, bias_list: str, num_layers: int) -> Tuple[nn.Module]:
"""
Constructs a dense neural network model.
Args:
dim_list (str): A comma-separated string representing layer dimensions.
bias_list (str): A comma-separated string representing bias usage (True/False) per layer.
num_layers (int): The number of layers in the model.
Returns:
Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
Raises:
ValueError: If `dim_list` length does not match `num_layers + 1`.
ValueError: If `bias_list` length does not match `num_layers`.
"""
dims = ast.literal_eval(dim_list)
biases = ast.literal_eval(bias_list)
if len(dims) - 1 != num_layers:
raise ValueError(
"dim_list and num_layers do not match. dim_list should contain "
"the input dimension for the first layer and output dimensions for all layers. "
"This list should be one element longer than the number of layers."
)
if len(biases) != num_layers:
raise ValueError("bias_list and num_layers do not match.")
with torch.inference_mode(False):
model = DenseModel(dims, biases, num_layers)
return (model,)
modules/pytorch_wrapper/ptn_linear_model_def.py
+99
View File
@@ -0,0 +1,99 @@
from typing import List, Dict, Any
import torch
import torch.nn as nn
class DenseSingleLayer(nn.Module):
"""A single dense layer followed by a ReLU activation.
Args:
in_features (int): Number of input features.
out_features (int): Number of output features.
bias (bool): If True, adds a learnable bias to the layer.
pragma: skip_doc
"""
@classmethod # TODO: Dummy method to suppress ComfyUI warning
def INPUT_TYPES(cls) -> Dict[str, Any]:
return {
}
def __init__(self, in_features: int, out_features: int, bias: bool) -> None:
super().__init__()
self.fc = nn.Linear(in_features, out_features, bias)
self.activation = nn.ReLU()
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
"""Forward pass of the dense layer with ReLU activation.
Args:
inputs (torch.Tensor): Input tensor of shape `(batch_size, in_features)`.
Returns:
torch.Tensor: Output tensor of shape `(batch_size, out_features)`.
"""
x = self.fc(inputs)
x = self.activation(x)
return x
class DenseModel(nn.Module):
"""A multi-layer dense neural network with ReLU activation after each dense layer, except for the final layer.
Args:
dim_list (List[int]): List of feature dimensions. Should be of length `num_layers + 1`.
bias_list (List[bool]): List indicating whether each layer should have a bias.
num_layers (int): Number of layers in the model.
Example:
For a three-layer network:
- Layer 1: input -> output1
- Layer 2: output1 -> output2
- Layer 3: output2 -> output3
`dim_list = [input_dim, output1_dim, output2_dim, output3_dim]`
pragma: skip_doc
"""
@classmethod # TODO: Dummy method to suppress ComfyUI warning
def INPUT_TYPES(cls) -> Dict[str, Any]:
return {
}
def __init__(self, dim_list: List[int], bias_list: List[bool], num_layers: int) -> None:
super().__init__()
self.no_activation_layers = nn.ModuleList([
DenseSingleLayer(dim_list[i], dim_list[i + 1], bias_list[i])
for i in range(num_layers - 1)
])
self.dim_list = dim_list
self.bias_list = bias_list
self.num_layers = num_layers
self.final_dense = nn.Linear(
dim_list[num_layers - 1], # in_features
dim_list[num_layers], # out_features
bias_list[num_layers - 1]
)
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
"""Forward pass of the dense model.
Args:
inputs (torch.Tensor): Input tensor of shape `(batch_size, dim_list[0])`.
Returns:
torch.Tensor: Output tensor of shape `(batch_size, dim_list[-1])`.
"""
x = inputs.view(-1, self.dim_list[0]) # Ensure correct input shape
for i in range(self.num_layers - 1):
x = self.no_activation_layers[i](x)
x = self.final_dense(x)
return x
modules/pytorch_wrapper/pto_adam.py
+46
View File
@@ -0,0 +1,46 @@
from typing import Any, Dict
import torch
import torch.nn as nn
import torch.nn.functional as F
class PtoAdam:
"""
Converts PyTorch tensor to NumPy ndarray.
category: PyTorch wrapper - Tensor data conversion
"""
@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": {
"model": ("PTMODEL", {}),
"learning_rate": ("FLOAT", {"default":0.001, "min":1e-10, "max":1, "step":0.0000000001}),
"beta1": ("FLOAT", {"default":0.9, "min":1e-10, "max":1, "step":0.0000000001}),
"beta2": ("FLOAT", {"default":0.999, "min":1e-10, "max":1, "step":0.0000000001}),
}
}
RETURN_TYPES: tuple = ("PTOPTIMIZER",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, model, learning_rate, beta1, beta2) -> tuple:
"""
Converts PyTorch tensor to NumPy ndarray.
Args:
tens (torch.Tensor): PyTorch Tensor
Returns:
tuple: A tuple containing the ndarray.
"""
with torch.inference_mode(False):
opt = torch.optim.Adam(model.parameters(), lr=learning_rate, betas=(beta1, beta2))
return(opt,)
modules/pytorch_wrapper/ptv_dataset.py
+124
View File
@@ -0,0 +1,124 @@
import inspect
from typing import Any, Dict, Optional, Callable, Tuple
import torch
import torchvision.datasets as datasets
import ast
from .utils import get_dataset_full_path
class PtvDataset:
"""
A Torchvision Dataset class wrapper.
Parameters
**name**
Specify a dataset class name such as MNIST, FashionMNIST in string.
This will be converted to a dataset class internally.
**download**
Set to True if you want to download.
**root**
Specify the root directory of the downloaded dataset relative to the
dataset directory of this extension, or specify the absolute path.
If you leave as a blank, the dataset will be downloaded under the `datasets` directory
of this extension.
**transform**
Data transforms e.g. transform to PyTorch tensor, normalize image.
Plug in the PTV Transforms node that contains Torchvision Transforms functionality.
**parameters**
Specify other parameters in Python dict format.
For example, if you want to set Train parameter to True to download the train set, specify:
```
{"train": True}
```
or if you want to download the test set, specify:
```
{"train": False}
```
category: PyTorch wrapper - Training
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"name": ("STRING", {"default": "", "multiline": False}),
"download": ("BOOLEAN", {"default": True}),
"root": ("STRING", {"default": "", "multiline": False}),
},
"optional": {
"transform": ("PTVTRANSFORM", {}),
"parameters": ("STRING", {"default": "", "multiline": True})
}
}
RETURN_TYPES: Tuple[str] = ("PTVDATASET",)
FUNCTION: str = "f"
CATEGORY: str = "Training"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self,
name: str,
download: bool,
root: str,
parameters: str,
transform: Optional[Callable] = None) -> Tuple:
"""
Loads a dataset from Torchvision with specified parameters.
Args:
name (str): Name of the dataset class (e.g., MNIST, FashionMNIST).
download (bool): Whether to download the dataset.
root (str): Root directory for storing the dataset.
parameters (str): Additional dataset parameters in stringified dictionary format.
transform (Optional[Callable]): Transformations to apply to the dataset.
Returns:
Tuple: A tuple containing the dataset instance.
"""
with torch.inference_mode(False):
param_dict = ast.literal_eval(parameters)
# Converts the datasets class name to dataset class itself
try:
dataset_class = getattr(datasets, name)
except AttributeError:
raise ValueError(f"Dataset '{name}' is not available in torchvision.datasets")
# Get valid parameters for the dataset class
valid_params = inspect.signature(dataset_class).parameters
# Filter only parameters that match the dataset's constructor
filtered_params = {k: v for k, v in param_dict.items() if k in valid_params}
if root:
dataset_path = get_dataset_full_path(root)
else: # Use default directory
dataset_path = get_dataset_full_path("")
if transform:
dc = dataset_class(download=download,
root=dataset_path,
transform=transform,
**filtered_params)
else:
dc = dataset_class(download=download,
root=dataset_path,
**filtered_params)
return (dc,)
modules/pytorch_wrapper/ptv_dataset_loader.py
+148
View File
@@ -0,0 +1,148 @@
import ast
import inspect
from typing import Any, Dict, Tuple
import torch
import torchvision.datasets as datasets
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from .utils import get_dataset_full_path
class PtvDatasetLoader:
"""
A node to combine the dataset and data loader into a single node.
Transform is internally called to convert input data to PyTorch tensors.
Parameters
**name**
Specify the dataset class name such as MNIST, FashionMNIST in string.
This will be converted to a dataset class internally.
**download**
Set to True if you want to download.
**root**
Specify the root directory of the downloaded dataset relative to the
dataset directory of this extension, or specify the absolute path.
If you leave as a blank, the dataset will be downloaded under the `datasets` directory
of this extension.
**batch_size**
The number of samples per batch.
**shuffle**
If True, shuffles the dataset before each epoch.
**dataset parameters**
Specify other parameters in Python dict format.
For example, if you want to set Train parameter to True to download the train set, specify:
```
{"train": True}
```
or if you want to download the test set, specify:
```
{"train": False}
```
**load parameters**
Additional parameters in Python dictionary format.
category: PyTorch wrapper - Training
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"name": ("STRING", {"default": "", "multiline": False}),
"download": ("BOOLEAN", {"default": True}),
"root": ("STRING", {"default": "", "multiline": False}),
"batch_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
"shuffle": ("BOOLEAN", {"default": False}),
},
"optional": {
"dataset_parameters": ("STRING", {"default": '{"train":True}', "multiline": True}),
"load_parameters": ("STRING", {"default": '{"num_workers:1}', "multiline": True})
}
}
RETURN_TYPES: Tuple[str] = ("PTDATALOADER",)
FUNCTION: str = "f"
CATEGORY: str = "Training"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self,
name: str,
download: bool,
root: str,
batch_size: int,
shuffle: bool,
dataset_parameters: str,
load_parameters: str) -> Tuple:
"""
Loads a dataset from Torchvision with specified parameters.
Args:
name (str): Name of the dataset class (e.g., MNIST, FashionMNIST).
download (bool): Whether to download the dataset.
root (str): Root directory for storing the dataset.
batch_size (int): The number of samples per batch.
shuffle (bool): If True, shuffles the dataset before each epoch.
dataset_parameters (str): Additional dataset parameters in stringified dictionary format.
load_parameters (str): Additional parameters for data load in Python dictionary format.
Returns:
Tuple: A tuple containing the dataset instance.
"""
with torch.inference_mode(False):
dataset_param_dict = ast.literal_eval(dataset_parameters)
# Converts the datasets class name to dataset class itself
try:
dataset_class = getattr(datasets, name)
except AttributeError:
raise ValueError(f"Dataset '{name}' is not available in torchvision.datasets")
# Get valid parameters for the dataset class
valid_params = inspect.signature(dataset_class).parameters
# Filter only parameters that match the dataset's constructor
filtered_params = {k: v for k, v in dataset_param_dict.items() if k in valid_params}
if root:
dataset_path = get_dataset_full_path(root)
else: # Use default directory
dataset_path = get_dataset_full_path("")
transform = transforms.Compose([transforms.ToTensor()])
dc = dataset_class(download=download,
root=dataset_path,
transform=transform,
**filtered_params)
load_param_dict = ast.literal_eval(load_parameters)
# Get valid parameters for the DataLoader class
valid_params = inspect.signature(DataLoader).parameters
# Filter only parameters that match the dataset's constructor
filtered_params = {k: v for k, v in load_param_dict.items() if k in valid_params}
dl = DataLoader(dataset=dc,
batch_size=batch_size,
shuffle=shuffle,
**filtered_params)
return (dl,)
modules/pytorch_wrapper/ptv_transforms_to_tensor.py
+37
View File
@@ -0,0 +1,37 @@
from typing import Any, Dict, Optional, Callable
import torchvision.transforms as transforms
class PtvTransformsToTensor:
"""
Transforms elements of dataset to PyTorch tensors.
category: PyTorch wrapper - Training
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: An empty dictionary.
"""
return {
}
RETURN_TYPES: tuple = ("PTVTRANSFORM",)
FUNCTION: str = "f"
CATEGORY: str = "Training"
def f(self) -> tuple:
"""
Transforms elements of dataset to PyTorch tensors.
Returns:
tuple: A tuple containing the transform callable.
"""
c = transforms.Compose([transforms.ToTensor()])
return(c,)
modules/pytorch_wrapper/utils.py
+54
View File
@@ -1,3 +1,4 @@
import os
import ast
from typing import Optional, Tuple, Union
import torch
@@ -16,6 +17,59 @@ DTYPE_MAPPING = {
}
def get_resource_full_path(resource_name: str, resource_path: str) -> str:
"""
Resolves the absolute path to a resource file within the "resources" directory and
ensures that the directory exists.
Args:
resource_path (str): The path to the resource file. If the path is a relative path,
the path is considered relative to the `resource` directory of this extension.
Returns:
str: The absolute path to the resource file.
"""
if os.path.isabs(resource_path):
resource_full_path = resource_path
else:
resource_full_path = os.path.realpath(
os.path.join(os.path.dirname(__file__), "..", "..", resource_name, resource_path))
resource_dir = os.path.dirname(resource_full_path)
if os.path.exists(resource_dir) is False:
os.makedirs(resource_dir)
return resource_full_path
def get_model_full_path(model_path: str) -> str:
"""
Resolves the absolute path to a model file within the "models" directory and
ensures that the directory exists.
Args:
model_path (str): The path to the model file. If the path is a relative path,
the path is considered relative to the `models` directory of this extension.
Returns:
str: The absolute path to the model file.
"""
return get_resource_full_path("models", model_path)
def get_dataset_full_path(dataset_path: str) -> str:
"""
Resolves the absolute path to the datasete within the "datasets" directory and
ensures that the directory exists.
Args:
dataset_path (str): The path to the dataset file. If the path is a relative path,
the path is considered relative to the `datasets` directory of this extension.
Returns:
str: The absolute path to the dataset file.
"""
return get_resource_full_path("datasets", dataset_path)
def str_to_dim(dim: str) -> Optional[Union[int, Tuple[int, ...]]]:
"""
Converts a string representation of a dimension into an integer or a tuple of integers.
pyproject.toml
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "data-analysis"
description = "Data analysis custom modules for ComfyUI - Use Pandas & Matplotlib within ComfyUI"
version = "1.4.6"
version = "2.0.0"
license = {file = "LICENSE"}
[project.urls]
requirements.txt
+2 -1
View File
@@ -1,4 +1,5 @@
matplotlib>=3.10.0
pandas>=2.2.3
seaborn>=0.13.2
scipy>=1.14.1
scipy>=1.14.1
scikit-learn>=1.6.1
tests/pytorch_wrapper/pt_data_loader_test.py
+43
View File
@@ -0,0 +1,43 @@
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.ptv_dataset import PtvDataset
from pytorch_wrapper.ptv_transforms_to_tensor import PtvTransformsToTensor
from pytorch_wrapper.pt_data_loader import PtDataLoader
class TestPtDataLoader(unittest.TestCase):
def setUp(self):
"""Set up test data."""
self.dataset_node = PtvDataset()
self.compose_node = PtvTransformsToTensor()
self.data_loader_node = PtDataLoader()
def test_with_transform(self):
"""Tests with transform"""
bs = 2
transform = self.compose_node.f()[0]
dataset = self.dataset_node.f("MNIST",
True, # download
"", # root
'{"train": True}',
transform=transform)[0]
train_data = self.data_loader_node.f(
dataset=dataset,
batch_size=bs,
shuffle=True,
parameters='{"num_workers":1}'
)[0]
x, y = next(iter(train_data))
self.assertEqual(x.size(), torch.Size([bs, 1, 28, 28]))
self.assertEqual(y.size(), torch.Size([bs]))
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_evaluate_classification_model_test.py
+79
View File
@@ -0,0 +1,79 @@
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_linear_model import PtnLinearModel
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 = PtnLinearModel()
self.model = self.model_node.f(
"784,180,42,10", # input dim for first layer, output dim for each layer
"True,True,True", # bias
3 # num layers
)[0]
# Load weights
self.load_model = PtLoadModel()
epochs = 1
self.loaded_model = self.load_model.f(self.model, f"fashion_mnist_{epochs}_epochs.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_model_test.py
+59
View File
@@ -0,0 +1,59 @@
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_linear_model import PtnLinearModel
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 TestPtTrainModel(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 = PtnLinearModel()
self.model = self.model_node.f(
"784,180,42,10", # input dim for first layer, output dim for each layer
"True,True,True", # bias
3 # num layers
)[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.pt")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_train_classification_model_with_dataset_loader_test.py
+51
View File
@@ -0,0 +1,51 @@
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_linear_model import PtnLinearModel
from pytorch_wrapper.pto_adam import PtoAdam
from pytorch_wrapper.ptv_dataset_loader import PtvDatasetLoader
from pytorch_wrapper.pt_train_classification_model import PtTrainClassificationModel
from pytorch_wrapper.pt_save_model import PtSaveModel
class TestPtTrainModelWithDatasetLoader(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.dataset_loader_node = PtvDatasetLoader()
self.train_loader = self.dataset_loader_node.f(
"MNIST",
True, # download
"", # root
batch_size=128,
shuffle=True,
dataset_parameters='{"train": True}',
load_parameters='{"num_workers":1}'
)[0]
self.model_node = PtnLinearModel()
self.model = self.model_node.f(
"784,180,42,10", # input dim for first layer, output dim for each layer
"True,True,True", # bias
3 # num layers
)[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 = 2
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"mnist_{epochs}_epochs.pt")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/ptn_linear_model_test.py
+29
View File
@@ -0,0 +1,29 @@
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_linear_model import PtnLinearModel # Updated import to PtnLinear
class TestPtnLinearModel(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.model_node = PtnLinearModel()
self.model = self.model_node.f(
"784,180,42,10", # input dim for first layer, output dim for each layer
"True,True,True", # bias
3 # num layers
)[0] # HINADA note: Using roughly 0.23x ratio for reduction
def test_forward(self):
x = torch.ones(4, 28, 28)
out = self.model(x)
self.assertEqual(out.size(), torch.Size([4, 10]))
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/ptv_dataset_loader_test.py
+34
View File
@@ -0,0 +1,34 @@
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.ptv_dataset_loader import PtvDatasetLoader
class TestPtDataLoader(unittest.TestCase):
def setUp(self):
"""Set up test data."""
self.dataset_loader_node = PtvDatasetLoader()
def test1(self):
"""Tests downloading data and checks dimension"""
bs = 2
train_data = self.dataset_loader_node.f(
"MNIST",
True, # download
"", # root
batch_size=bs,
shuffle=True,
dataset_parameters='{"train": True}',
load_parameters='{"num_workers":1}'
)[0]
x, y = next(iter(train_data))
self.assertEqual(x.size(), torch.Size([bs, 1, 28, 28]))
self.assertEqual(y.size(), torch.Size([bs]))
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/ptv_dataset_test.py
+36
View File
@@ -0,0 +1,36 @@
import os
import sys
import unittest
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.ptv_dataset import PtvDataset
from pytorch_wrapper.ptv_transforms_to_tensor import PtvTransformsToTensor
class TestPtvDataset(unittest.TestCase):
def setUp(self):
"""Set up test data."""
self.node = PtvDataset()
self.compose_node = PtvTransformsToTensor()
def test_no_transform(self):
"""Tests downloading without transform"""
self.node.f("MNIST",
True, # download
"", # root
'{"train": True}',
transform=None)
def test_with_transform(self):
"""Tests downloading with transform"""
transform = self.compose_node.f()[0]
self.node.f("MNIST",
True, # download
"", # root
'{"train": False}',
transform=transform)
if __name__ == "__main__":
unittest.main()