diff --git a/.gitignore b/.gitignore index cea1d96..17a60fc 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ __pycache__/ **/__pycache__/ - +datasets/ +models/ +*.pth diff --git a/README.md b/README.md index f5ca4ff..0e47476 100644 --- a/README.md +++ b/README.md @@ -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. + +
+ ![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 ``` diff --git a/__init__.py b/__init__.py index 3a596ea..d070694 100644 --- a/__init__.py +++ b/__init__.py @@ -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", diff --git a/docs/images/fashion_mnist_eval.png b/docs/images/fashion_mnist_eval.png new file mode 100644 index 0000000..8278581 Binary files /dev/null and b/docs/images/fashion_mnist_eval.png differ diff --git a/docs/images/fashion_mnist_train.png b/docs/images/fashion_mnist_train.png new file mode 100644 index 0000000..c74a788 Binary files /dev/null and b/docs/images/fashion_mnist_train.png differ diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index 285b3fe..c6e6374 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -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 | | --- | --- | diff --git a/docs/reference/pt_data_loader.md b/docs/reference/pt_data_loader.md new file mode 100644 index 0000000..2ae51c8 --- /dev/null +++ b/docs/reference/pt_data_loader.md @@ -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 | + +
+Category: PyTorch wrapper - Training + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/pt_evaluate_classification_model.md b/docs/reference/pt_evaluate_classification_model.md new file mode 100644 index 0000000..c78bce9 --- /dev/null +++ b/docs/reference/pt_evaluate_classification_model.md @@ -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 | + +
+Category: PyTorch wrapper - Training + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/pt_load_model.md b/docs/reference/pt_load_model.md new file mode 100644 index 0000000..0466004 --- /dev/null +++ b/docs/reference/pt_load_model.md @@ -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 | + +
+Category: PyTorch wrapper - Training + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/pt_save_model.md b/docs/reference/pt_save_model.md new file mode 100644 index 0000000..da2d30c --- /dev/null +++ b/docs/reference/pt_save_model.md @@ -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 | + +
+Category: IO + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/pt_train_classification_model.md b/docs/reference/pt_train_classification_model.md new file mode 100644 index 0000000..1ceaac6 --- /dev/null +++ b/docs/reference/pt_train_classification_model.md @@ -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 | + +
+Category: PyTorch wrapper - Training + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptn_linear_model.md b/docs/reference/ptn_linear_model.md new file mode 100644 index 0000000..a13e46c --- /dev/null +++ b/docs/reference/ptn_linear_model.md @@ -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 | + +
+Category: PyTorch wrapper - Training + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptn_linear_model_def.md b/docs/reference/ptn_linear_model_def.md new file mode 100644 index 0000000..1841a98 --- /dev/null +++ b/docs/reference/ptn_linear_model_def.md @@ -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]` + +
+Category: Miscellaneous + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/pto_adam.md b/docs/reference/pto_adam.md new file mode 100644 index 0000000..fa394e9 --- /dev/null +++ b/docs/reference/pto_adam.md @@ -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 | + +
+Category: PyTorch wrapper - Tensor data conversion + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptv_dataset.md b/docs/reference/ptv_dataset.md new file mode 100644 index 0000000..1258c76 --- /dev/null +++ b/docs/reference/ptv_dataset.md @@ -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 | + +
+Category: PyTorch wrapper - Training + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptv_dataset_loader.md b/docs/reference/ptv_dataset_loader.md new file mode 100644 index 0000000..ee13ac5 --- /dev/null +++ b/docs/reference/ptv_dataset_loader.md @@ -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 | + +
+Category: PyTorch wrapper - Training + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptv_transforms_to_tensor.md b/docs/reference/ptv_transforms_to_tensor.md new file mode 100644 index 0000000..e1ca095 --- /dev/null +++ b/docs/reference/ptv_transforms_to_tensor.md @@ -0,0 +1,12 @@ +# Ptv Transforms To Tensor +Transforms elements of dataset to PyTorch tensors. + +## Output +| Data type | +|---| +| Ptvtransform | + +
+Category: PyTorch wrapper - Training + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/examples/workflows/fashion_mnist_eval.json b/examples/workflows/fashion_mnist_eval.json new file mode 100644 index 0000000..c9e452e --- /dev/null +++ b/examples/workflows/fashion_mnist_eval.json @@ -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} \ No newline at end of file diff --git a/examples/workflows/fashion_mnist_train.json b/examples/workflows/fashion_mnist_train.json new file mode 100644 index 0000000..c693a15 --- /dev/null +++ b/examples/workflows/fashion_mnist_train.json @@ -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} \ No newline at end of file diff --git a/modules/pytorch_wrapper/pt_data_loader.py b/modules/pytorch_wrapper/pt_data_loader.py new file mode 100644 index 0000000..7ec163b --- /dev/null +++ b/modules/pytorch_wrapper/pt_data_loader.py @@ -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,) diff --git a/modules/pytorch_wrapper/pt_evaluate_classification_model.py b/modules/pytorch_wrapper/pt_evaluate_classification_model.py new file mode 100644 index 0000000..92fb68a --- /dev/null +++ b/modules/pytorch_wrapper/pt_evaluate_classification_model.py @@ -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) diff --git a/modules/pytorch_wrapper/pt_load_model.py b/modules/pytorch_wrapper/pt_load_model.py new file mode 100644 index 0000000..c40dbbc --- /dev/null +++ b/modules/pytorch_wrapper/pt_load_model.py @@ -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,) \ No newline at end of file diff --git a/modules/pytorch_wrapper/pt_save_model.py b/modules/pytorch_wrapper/pt_save_model.py new file mode 100644 index 0000000..993c221 --- /dev/null +++ b/modules/pytorch_wrapper/pt_save_model.py @@ -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),) \ No newline at end of file diff --git a/modules/pytorch_wrapper/pt_train_classification_model.py b/modules/pytorch_wrapper/pt_train_classification_model.py new file mode 100644 index 0000000..632d0be --- /dev/null +++ b/modules/pytorch_wrapper/pt_train_classification_model.py @@ -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))) \ No newline at end of file diff --git a/modules/pytorch_wrapper/ptn_linear_model.py b/modules/pytorch_wrapper/ptn_linear_model.py new file mode 100644 index 0000000..3f1c9f5 --- /dev/null +++ b/modules/pytorch_wrapper/ptn_linear_model.py @@ -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,) diff --git a/modules/pytorch_wrapper/ptn_linear_model_def.py b/modules/pytorch_wrapper/ptn_linear_model_def.py new file mode 100644 index 0000000..28a10e6 --- /dev/null +++ b/modules/pytorch_wrapper/ptn_linear_model_def.py @@ -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 diff --git a/modules/pytorch_wrapper/pto_adam.py b/modules/pytorch_wrapper/pto_adam.py new file mode 100644 index 0000000..36fcef7 --- /dev/null +++ b/modules/pytorch_wrapper/pto_adam.py @@ -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,) \ No newline at end of file diff --git a/modules/pytorch_wrapper/ptv_dataset.py b/modules/pytorch_wrapper/ptv_dataset.py new file mode 100644 index 0000000..29984fd --- /dev/null +++ b/modules/pytorch_wrapper/ptv_dataset.py @@ -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,) diff --git a/modules/pytorch_wrapper/ptv_dataset_loader.py b/modules/pytorch_wrapper/ptv_dataset_loader.py new file mode 100644 index 0000000..40ad553 --- /dev/null +++ b/modules/pytorch_wrapper/ptv_dataset_loader.py @@ -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,) \ No newline at end of file diff --git a/modules/pytorch_wrapper/ptv_transforms_to_tensor.py b/modules/pytorch_wrapper/ptv_transforms_to_tensor.py new file mode 100644 index 0000000..89fdded --- /dev/null +++ b/modules/pytorch_wrapper/ptv_transforms_to_tensor.py @@ -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,) + + diff --git a/modules/pytorch_wrapper/utils.py b/modules/pytorch_wrapper/utils.py index 4b459ea..f9e9cb6 100644 --- a/modules/pytorch_wrapper/utils.py +++ b/modules/pytorch_wrapper/utils.py @@ -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. diff --git a/pyproject.toml b/pyproject.toml index 01503b6..600e97f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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] diff --git a/requirements.txt b/requirements.txt index 2dd4374..58e2f64 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ matplotlib>=3.10.0 pandas>=2.2.3 seaborn>=0.13.2 -scipy>=1.14.1 \ No newline at end of file +scipy>=1.14.1 +scikit-learn>=1.6.1 \ No newline at end of file diff --git a/tests/pytorch_wrapper/pt_data_loader_test.py b/tests/pytorch_wrapper/pt_data_loader_test.py new file mode 100644 index 0000000..6b59a75 --- /dev/null +++ b/tests/pytorch_wrapper/pt_data_loader_test.py @@ -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() diff --git a/tests/pytorch_wrapper/pt_evaluate_classification_model_test.py b/tests/pytorch_wrapper/pt_evaluate_classification_model_test.py new file mode 100644 index 0000000..262e13a --- /dev/null +++ b/tests/pytorch_wrapper/pt_evaluate_classification_model_test.py @@ -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() diff --git a/tests/pytorch_wrapper/pt_train_classification_model_test.py b/tests/pytorch_wrapper/pt_train_classification_model_test.py new file mode 100644 index 0000000..9edb822 --- /dev/null +++ b/tests/pytorch_wrapper/pt_train_classification_model_test.py @@ -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() diff --git a/tests/pytorch_wrapper/pt_train_classification_model_with_dataset_loader_test.py b/tests/pytorch_wrapper/pt_train_classification_model_with_dataset_loader_test.py new file mode 100644 index 0000000..43fead0 --- /dev/null +++ b/tests/pytorch_wrapper/pt_train_classification_model_with_dataset_loader_test.py @@ -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() diff --git a/tests/pytorch_wrapper/ptn_linear_model_test.py b/tests/pytorch_wrapper/ptn_linear_model_test.py new file mode 100644 index 0000000..124a15a --- /dev/null +++ b/tests/pytorch_wrapper/ptn_linear_model_test.py @@ -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() diff --git a/tests/pytorch_wrapper/ptv_dataset_loader_test.py b/tests/pytorch_wrapper/ptv_dataset_loader_test.py new file mode 100644 index 0000000..0d01690 --- /dev/null +++ b/tests/pytorch_wrapper/ptv_dataset_loader_test.py @@ -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() diff --git a/tests/pytorch_wrapper/ptv_dataset_test.py b/tests/pytorch_wrapper/ptv_dataset_test.py new file mode 100644 index 0000000..3eeffba --- /dev/null +++ b/tests/pytorch_wrapper/ptv_dataset_test.py @@ -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()