diff --git a/README.md b/README.md index c9c0469..b2ba28c 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ Check out [this tutorial on classifying dogs and cats](docs/dog_cat_classificati For visualization or analysis features, install ComfyUI-Data-Analysis as well. -The workflow directory contains an easy-to-use workflow for training a high-performing ResNet model on CIFAR-10. +The workflow directory offers user-friendly example workflows for training a high-performing ResNet model on CIFAR-10. One example achieved **94.34% accuracy** on the validation set. While not a state-of-the-art metric, it provides a solid starting point for customizing the workflow to your needs. **ResNet model** ![Train](docs/images/resnet_train.png) diff --git a/__init__.py b/__init__.py index a32bf64..ae33279 100644 --- a/__init__.py +++ b/__init__.py @@ -93,6 +93,7 @@ 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_train_classification_model_lr import PtTrainClassificationModelLr from .modules.pytorch_wrapper.pt_unsqueeze import PtUnsqueeze from .modules.pytorch_wrapper.pt_var import PtVar from .modules.pytorch_wrapper.pt_view import PtView @@ -102,6 +103,9 @@ from .modules.pytorch_wrapper.ptn_conv_model import PtnConvModel from .modules.pytorch_wrapper.ptn_linear_model import PtnLinearModel from .modules.pytorch_wrapper.ptn_resnet_model import PtnResnetModel from .modules.pytorch_wrapper.pto_adam import PtoAdam +from .modules.pytorch_wrapper.pto_lr_scheduler_cosine_annealing import PtoLrSchedulerCosineAnnealing +from .modules.pytorch_wrapper.pto_lr_scheduler_reduce_on_plateau import PtoLrSchedulerReduceOnPlateau +from .modules.pytorch_wrapper.pto_lr_scheduler_step import PtoLrSchedulerStep from .modules.pytorch_wrapper.ptv_dataset import PtvDataset from .modules.pytorch_wrapper.ptv_dataset_len import PtvDatasetLen from .modules.pytorch_wrapper.ptv_dataset_loader import PtvDatasetLoader @@ -211,6 +215,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtToNumpy": PtToNumpy, "PtToRgbTensors": PtToRgbTensors, "PtTrainClassificationModel": PtTrainClassificationModel, + "PtTrainClassificationModelLr": PtTrainClassificationModelLr, "PtUnsqueeze": PtUnsqueeze, "PtVar": PtVar, "PtView": PtView, @@ -220,6 +225,9 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtnLinearModel": PtnLinearModel, "PtnResnetModel": PtnResnetModel, "PtoAdam": PtoAdam, + "PtoLrSchedulerCosineAnnealing": PtoLrSchedulerCosineAnnealing, + "PtoLrSchedulerReduceOnPlateau": PtoLrSchedulerReduceOnPlateau, + "PtoLrSchedulerStep": PtoLrSchedulerStep, "PtvDataset": PtvDataset, "PtvDatasetLen": PtvDatasetLen, "PtvDatasetLoader": PtvDatasetLoader, @@ -329,6 +337,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtToNumpy": "Pt To Numpy", "PtToRgbTensors": "Pt To Rgb Tensors", "PtTrainClassificationModel": "Pt Train Classification Model", + "PtTrainClassificationModelLr": "Pt Train Classification Model Lr", "PtUnsqueeze": "Pt Unsqueeze", "PtVar": "Pt Var", "PtView": "Pt View", @@ -338,6 +347,9 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtnLinearModel": "Ptn Linear Model", "PtnResnetModel": "Ptn Resnet Model", "PtoAdam": "Pto Adam", + "PtoLrSchedulerCosineAnnealing": "Pto Lr Scheduler Cosine Annealing", + "PtoLrSchedulerReduceOnPlateau": "Pto Lr Scheduler Reduce On Plateau", + "PtoLrSchedulerStep": "Pto Lr Scheduler Step", "PtvDataset": "Ptv Dataset", "PtvDatasetLen": "Ptv Dataset Len", "PtvDatasetLoader": "Ptv Dataset Loader", diff --git a/docs/model_training.md b/docs/model_training.md index 9d65c19..545e63b 100644 --- a/docs/model_training.md +++ b/docs/model_training.md @@ -35,10 +35,15 @@ Convolutional model * Evaluation: examples/workflows/fashion_mnist_eval_conv.json ResNet model -* Training: examples/workflows/cifar10_train_v8.json -* Evaluation: examples/workflows/cifar10_eval_v8.json +**Training** +* examples/workflows/cifar10_train_v8.json +* examples/workflows/cifar10_train_v11_2nd_run.json + +**Evaluation** +* examples/workflows/cifar10_eval_v8.json +* examples/workflows/cifar10_eval_v11_2nd_run.json -This model achieved 93.12% accuracy on CIFAR-10. While not state-of-the-art, it is a strong result. Due to random initialization, data augmentation, and dataset shuffling, accuracy may vary slightly between runs. If your accuracy is significantly lower, consider running it multiple times or checking for potential issues in the training setup. +This train_v8 model achieved 93.12% accuracy & train_v11_2nd_run model achieved 94.34 on CIFAR-10 (first run was 94.27%). train_v11_2nd_run added cosine annealing learning rate scheduler while train_v8 used fixed learning rate. While not state-of-the-art, these are strong results. Due to random initialization, data augmentation, and dataset shuffling, accuracy may vary slightly between runs. If your accuracy is significantly lower, consider running it multiple times or checking for potential issues in the training setup. **Note: Before you try these workflows, please install ComfyUI-Data-Analysis extension.** diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index 407faac..189cc44 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -136,10 +136,14 @@ | [Pt Load Model](pt_load_model.md) | A wrapper class for saving a PyTorch model. | | [Pt Predict Classification Model](pt_predict_classification_model.md) | Performs inference on input data. | | [Pt Train Classification Model](pt_train_classification_model.md) | Trains a classification model using a given dataset, optimizer, and number of epochs. | +| [Pt Train Classification Model Lr](pt_train_classification_model_lr.md) | Trains a classification model using a given dataset, optimizer, and number of epochs with learning rate decay. | | [Ptn Conv Model](ptn_conv_model.md) | A convolutional model consisting of multiple convolutional layers. | | [Ptn Linear Model](ptn_linear_model.md) | A linear model consisting of dense layers. | | [Ptn Resnet Model](ptn_resnet_model.md) | A Resnet model consisting of multiple Resnet layers. | | [Pto Adam](pto_adam.md) | Instantiates the Adam optimizer. | +| [Pto Lr Scheduler Cosine Annealing](pto_lr_scheduler_cosine_annealing.md) | Creates a cosine annealing learning rate scheduler for an optimizer. | +| [Pto Lr Scheduler Reduce On Plateau](pto_lr_scheduler_reduce_on_plateau.md) | Creates a reduce-on-plateau learning rate scheduler for an optimizer. | +| [Pto Lr Scheduler Step](pto_lr_scheduler_step.md) | Creates a StepLR learning rate scheduler for an optimizer. | | [Ptv Dataset](ptv_dataset.md) | A Torchvision Dataset class wrapper. | | [Ptv Dataset Len](ptv_dataset_len.md) | A wrapper class that calls Python len on a dataset. | | [Ptv Dataset Loader](ptv_dataset_loader.md) | A node to combine the dataset and data loader into a single node. | diff --git a/docs/reference/pt_train_classification_model_lr.md b/docs/reference/pt_train_classification_model_lr.md new file mode 100644 index 0000000..a2e9d61 --- /dev/null +++ b/docs/reference/pt_train_classification_model_lr.md @@ -0,0 +1,31 @@ +# Pt Train Classification Model Lr +Trains a classification model using a given dataset, optimizer, and number of epochs with learning rate decay. + +early_stopping_rounds: Specifies the number of rounds to monitor the evaluation loss when early_stopping is enabled and validation data is specified. If the loss does not decrease compared to the current best loss within this period, training will be stopped. +output_best_val_model: Whether to outputs the model with the best val loss or the model from the last epoch. + +## Input +| Name | Data type | +|---|---| +| model | Ptmodel | +| train_loader | Ptdataloader | +| optimizer | Ptoptimizer | +| scheduler | Ptlrscheduler | +| epochs | Int | +| use_gpu | Boolean | +| early_stopping | Boolean | +| early_stopping_rounds | Int | +| output_best_val_model | 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/pto_lr_scheduler_cosine_annealing.md b/docs/reference/pto_lr_scheduler_cosine_annealing.md new file mode 100644 index 0000000..764b260 --- /dev/null +++ b/docs/reference/pto_lr_scheduler_cosine_annealing.md @@ -0,0 +1,49 @@ +# Pto Lr Scheduler Cosine Annealing +Creates a cosine annealing learning rate scheduler for an optimizer. + +https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.CosineAnnealingLR.html +lists: + +eta_t = eta_min + (1/2) * (eta_max - eta_min) * (1 + cos((t_cur / t_max) * pi)) + +eta_min is set to 0 by default, so the equation simplifies to: + +eta_t = (1/2) * eta_max * (1 + cos((t_cur / t_max) * pi)) + +When t_cur = 0: + (1 + cos((t_cur / t_max) * pi)) evaluates to (1 + cos(0 * pi)) = (1 + 1) = 2 + +When t_cur = t_max: + (1 + cos((t_cur / t_max) * pi)) evaluates to (1 + cos(1 * pi)) = (1 - 1) = 0 + +Therefore, eta_t value changes as follows: +* Starts from (1/2) * eta_max * 2, which simplifies to eta_max +* Ends at 0 + +Parameters + +**optimizer** + (torch.optim.Optimizer) – Wrapped optimizer whose learning rate needs scheduling. + +**num_epochs** + (int) – Number of epochs. This maps to the T_max parameter for CosineAnnealingLR internally. + +**minimum_lr** + (float) – Minimum learning rate. + +## Input +| Name | Data type | +|---|---| +| optimizer | Ptoptimizer | +| num_epochs | Int | +| minimum_lr | Float | + +## Output +| Data type | +|---| +| Ptlrscheduler | + +
+Category: PyTorch wrapper - Training + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/pto_lr_scheduler_reduce_on_plateau.md b/docs/reference/pto_lr_scheduler_reduce_on_plateau.md new file mode 100644 index 0000000..7d0e49f --- /dev/null +++ b/docs/reference/pto_lr_scheduler_reduce_on_plateau.md @@ -0,0 +1,30 @@ +# Pto Lr Scheduler Reduce On Plateau +Creates a reduce-on-plateau learning rate scheduler for an optimizer. + +Parameters + +**optimizer** + (torch.optim.Optimizer) – Wrapped optimizer whose learning rate needs scheduling. + +**grace_period** + (int) – Number of epochs to monitor the val loss reduction. If loss does not decrease within grace period epochs, learning rate will be reduced by multiplying the gamma. E.g. if gamma is set to 0.5, new learning rate will be 0.5 * current learning rate. + +**gamma** + (float) – Learning rate multiplier to be applied when grace period is up. + +## Input +| Name | Data type | +|---|---| +| optimizer | Ptoptimizer | +| grace_period | Int | +| gamma | Float | + +## Output +| Data type | +|---| +| Ptlrscheduler | + +
+Category: PyTorch wrapper - Training + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/pto_lr_scheduler_step.md b/docs/reference/pto_lr_scheduler_step.md new file mode 100644 index 0000000..f1eebfc --- /dev/null +++ b/docs/reference/pto_lr_scheduler_step.md @@ -0,0 +1,30 @@ +# Pto Lr Scheduler Step +Creates a StepLR learning rate scheduler for an optimizer. + +Parameters + +**optimizer** + (torch.optim.Optimizer) – Wrapped optimizer whose learning rate needs scheduling. + +**step_size** + (int) – Period of learning rate decay. + +**gamma** + (float) – Multiplicative factor of learning rate decay. + +## Input +| Name | Data type | +|---|---| +| optimizer | Ptoptimizer | +| step_size | Int | +| gamma | Float | + +## Output +| Data type | +|---| +| Ptlrscheduler | + +
+Category: PyTorch wrapper - Training + +ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/examples/workflows/cifar10_eval_v11_2nd_run.json b/examples/workflows/cifar10_eval_v11_2nd_run.json new file mode 100644 index 0000000..e2c84c8 --- /dev/null +++ b/examples/workflows/cifar10_eval_v11_2nd_run.json @@ -0,0 +1 @@ +{"last_node_id":92,"last_link_id":107,"nodes":[{"id":54,"type":"MarkdownNote","pos":[1712.6656494140625,2128.33642578125],"size":[210,58.94380569458008],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Load val data"],"color":"#432","bgcolor":"#653"},{"id":78,"type":"PtvDatasetLen","pos":[2045.1402587890625,2289.77490234375],"size":[210,26],"flags":{},"order":8,"mode":0,"inputs":[{"name":"dataset","type":"PTVDATASET","link":86}],"outputs":[{"name":"INT","type":"INT","links":[85],"slot_index":0}],"properties":{"Node name for S&R":"PtvDatasetLen"},"widgets_values":[]},{"id":79,"type":"CDAShowInt","pos":[2045.1402587890625,2366.537353515625],"size":[315,76],"flags":{},"order":10,"mode":0,"inputs":[{"name":"int_scalar","type":"INT","link":85,"widget":{"name":"int_scalar"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"properties":{"Node name for S&R":"CDAShowInt"},"widgets_values":[0,"10000"]},{"id":62,"type":"PtvDataset","pos":[2045.1402587890625,2050.991943359375],"size":[400,200],"flags":{},"order":5,"mode":0,"inputs":[{"name":"transform","type":"PTVTRANSFORM","link":89,"shape":7}],"outputs":[{"name":"PTVDATASET","type":"PTVDATASET","links":[75,86],"slot_index":0}],"properties":{"Node name for S&R":"PtvDataset"},"widgets_values":["CIFAR10",true,"","{\"train\":False}"]},{"id":84,"type":"CDAShowFloat","pos":[3364.6611328125,2185.21142578125],"size":[210,76],"flags":{},"order":13,"mode":0,"inputs":[{"name":"float_scalar","type":"FLOAT","link":95,"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.9434000000000001"]},{"id":85,"type":"CDAShowFloat","pos":[3637.660888671875,2183.1708984375],"size":[210,76],"flags":{},"order":14,"mode":0,"inputs":[{"name":"float_scalar","type":"FLOAT","link":96,"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.943420064094697"]},{"id":83,"type":"CDAShowInt","pos":[3364.6611328125,2323.3916015625],"size":[210,79.09933471679688],"flags":{},"order":15,"mode":0,"inputs":[{"name":"int_scalar","type":"INT","link":94,"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":81,"type":"PtvTransformsToTensor","pos":[1708.258544921875,2050.991943359375],"size":[302.4000244140625,26],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[{"name":"PTVTRANSFORM","type":"PTVTRANSFORM","links":[89],"slot_index":0}],"properties":{"Node name for S&R":"PtvTransformsToTensor"},"widgets_values":[]},{"id":86,"type":"CDAShowFloat","pos":[3637.660888671875,2050.991943359375],"size":[210,76],"flags":{},"order":12,"mode":0,"inputs":[{"name":"float_scalar","type":"FLOAT","link":97,"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.9436015123002528"]},{"id":88,"type":"PtLoadModel","pos":[2973.5009765625,1907.97412109375],"size":[303.63604736328125,75.06120300292969],"flags":{},"order":6,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":101},{"name":"model_path","type":"STRING","link":100,"widget":{"name":"model_path"}}],"outputs":[{"name":"model","type":"PTMODEL","links":[93],"slot_index":0}],"properties":{"Node name for S&R":"PtLoadModel"},"widgets_values":["/container/transfer/models/food101_conv_ep5_v2.pt"]},{"id":80,"type":"PtnResnetModel","pos":[2498.54833984375,1782.458740234375],"size":[315,106],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[101],"slot_index":0}],"properties":{"Node name for S&R":"PtnResnetModel"},"widgets_values":["(3,32,32)",10,10]},{"id":52,"type":"MarkdownNote","pos":[2498.54833984375,1657.6710205078125],"size":[210,58.94380569458008],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Define model"],"color":"#432","bgcolor":"#653"},{"id":64,"type":"PtDataLoader","pos":[2498.54833984375,2050.991943359375],"size":[313.175048828125,152.2939453125],"flags":{},"order":7,"mode":0,"inputs":[{"name":"dataset","type":"PTVDATASET","link":75}],"outputs":[{"name":"PTDATALOADER","type":"PTDATALOADER","links":[102],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[64,false,"{\"num_workers\":1}"]},{"id":82,"type":"PtEvaluateClassificationModel","pos":[2968.29638671875,2050.991943359375],"size":[313.0364990234375,146.49288940429688],"flags":{},"order":9,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":93},{"name":"data_loader","type":"PTDATALOADER","link":102}],"outputs":[{"name":"accuracy","type":"FLOAT","links":[107],"slot_index":0},{"name":"precision","type":"FLOAT","links":[97],"slot_index":1},{"name":"recall","type":"FLOAT","links":[95],"slot_index":2},{"name":"f1","type":"FLOAT","links":[96],"slot_index":3},{"name":"num_samples","type":"INT","links":[94],"slot_index":4}],"properties":{"Node name for S&R":"PtEvaluateClassificationModel"},"widgets_values":[true]},{"id":92,"type":"CDAShowFloat","pos":[3368.79443359375,2061.034423828125],"size":[210,76],"flags":{},"order":11,"mode":0,"inputs":[{"name":"float_scalar","type":"FLOAT","link":107,"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.9434"]},{"id":89,"type":"CDAStringCreate","pos":[2968.29638671875,1782.458740234375],"size":[315,58],"flags":{},"order":4,"mode":0,"inputs":[],"outputs":[{"name":"STRING","type":"STRING","links":[100],"slot_index":0}],"properties":{"Node name for S&R":"CDAStringCreate"},"widgets_values":["/container/transfer/models/cifar10_v11_2nd_run.pt"]}],"links":[[75,62,0,64,0,"PTVDATASET"],[85,78,0,79,0,"INT"],[86,62,0,78,0,"PTVDATASET"],[89,81,0,62,0,"PTVTRANSFORM"],[93,88,0,82,0,"PTMODEL"],[94,82,4,83,0,"INT"],[95,82,2,84,0,"FLOAT"],[96,82,3,85,0,"FLOAT"],[97,82,1,86,0,"FLOAT"],[100,89,0,88,1,"STRING"],[101,80,0,88,0,"PTMODEL"],[102,64,0,82,1,"PTDATALOADER"],[107,82,0,92,0,"FLOAT"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.9607449831913665,"offset":[-2381.931979888728,-1617.8240954143464]},"node_versions":{"pt-wrapper":"eb9d05ed533f01888758f204deea434badf4a0a9","data-analysis":"bde4a07755390554f858e0b8fd8554a3123a6612"},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/examples/workflows/cifar10_train_v11_2nd_run.json b/examples/workflows/cifar10_train_v11_2nd_run.json new file mode 100644 index 0000000..1fcafe1 --- /dev/null +++ b/examples/workflows/cifar10_train_v11_2nd_run.json @@ -0,0 +1 @@ +{"last_node_id":93,"last_link_id":112,"nodes":[{"id":61,"type":"PtvDataset","pos":[2007.994873046875,1723.90625],"size":[400,200],"flags":{},"order":9,"mode":0,"inputs":[{"name":"transform","type":"PTVTRANSFORM","link":103,"shape":7}],"outputs":[{"name":"PTVDATASET","type":"PTVDATASET","links":[74,83],"slot_index":0}],"properties":{"Node name for S&R":"PtvDataset"},"widgets_values":["CIFAR10",true,"","{\"train\":True}"]},{"id":19,"type":"PtSaveModel","pos":[3530.668701171875,1647.40185546875],"size":[396.8008117675781,75.65961456298828],"flags":{},"order":20,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":108},{"name":"model_path","type":"STRING","link":99,"widget":{"name":"model_path"}}],"outputs":[{"name":"model full path","type":"STRING","links":[70],"slot_index":0}],"properties":{"Node name for S&R":"PtSaveModel"},"widgets_values":["/container/transfer/models/cifar10.pt"]},{"id":62,"type":"PtvDataset","pos":[2007.994873046875,2039.8226318359375],"size":[400,200],"flags":{},"order":10,"mode":0,"inputs":[{"name":"transform","type":"PTVTRANSFORM","link":89,"shape":7}],"outputs":[{"name":"PTVDATASET","type":"PTVDATASET","links":[75,86],"slot_index":0}],"properties":{"Node name for S&R":"PtvDataset"},"widgets_values":["CIFAR10",true,"","{\"train\":False}"]},{"id":76,"type":"PtvDatasetLen","pos":[2007.994873046875,1638.4072265625],"size":[210,26],"flags":{},"order":13,"mode":0,"inputs":[{"name":"dataset","type":"PTVDATASET","link":83}],"outputs":[{"name":"INT","type":"INT","links":[84],"slot_index":0}],"properties":{"Node name for S&R":"PtvDatasetLen"},"widgets_values":[]},{"id":78,"type":"PtvDatasetLen","pos":[2007.994873046875,2297.724365234375],"size":[210,26],"flags":{},"order":15,"mode":0,"inputs":[{"name":"dataset","type":"PTVDATASET","link":86}],"outputs":[{"name":"INT","type":"INT","links":[85],"slot_index":0}],"properties":{"Node name for S&R":"PtvDatasetLen"},"widgets_values":[]},{"id":79,"type":"CDAShowInt","pos":[2007.994873046875,2381.43359375],"size":[315,76],"flags":{},"order":18,"mode":0,"inputs":[{"name":"int_scalar","type":"INT","link":85,"widget":{"name":"int_scalar"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"properties":{"Node name for S&R":"CDAShowInt"},"widgets_values":[0,"10000"]},{"id":53,"type":"MarkdownNote","pos":[1215.7994384765625,1723.90625],"size":[210,58.94380569458008],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Load train data"],"color":"#432","bgcolor":"#653"},{"id":54,"type":"MarkdownNote","pos":[1214.4312744140625,2039.8226318359375],"size":[210,58.94380569458008],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Load val data"],"color":"#432","bgcolor":"#653"},{"id":77,"type":"CDAShowInt","pos":[2007.994873046875,1504.9381103515625],"size":[315,76],"flags":{},"order":17,"mode":0,"inputs":[{"name":"int_scalar","type":"INT","link":84,"widget":{"name":"int_scalar"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"properties":{"Node name for S&R":"CDAShowInt"},"widgets_values":[0,"50000"]},{"id":90,"type":"PtvTransformsDataAugment","pos":[1595.48486328125,1723.90625],"size":[340.20001220703125,202],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[{"name":"PTVTRANSFORM","type":"PTVTRANSFORM","links":[103],"slot_index":0}],"properties":{"Node name for S&R":"PtvTransformsDataAugment"},"widgets_values":[0.5,0,0,0.125,0.125,1,1]},{"id":81,"type":"PtvTransformsToTensor","pos":[1595.48486328125,2039.8226318359375],"size":[338.8136291503906,26],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[{"name":"PTVTRANSFORM","type":"PTVTRANSFORM","links":[89],"slot_index":0}],"properties":{"Node name for S&R":"PtvTransformsToTensor"},"widgets_values":[]},{"id":56,"type":"MarkdownNote","pos":[3535.058349609375,1775.04931640625],"size":[210,58.94380569458008],"flags":{},"order":4,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Save model"],"color":"#432","bgcolor":"#653"},{"id":51,"type":"PandasShowText","pos":[3536.080078125,1878.28125],"size":[333.6095275878906,79.25810241699219],"flags":{},"order":21,"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":["","/container/transfer/models/cifar10_v11_2nd_run.pt"]},{"id":52,"type":"MarkdownNote","pos":[2485.625732421875,1294.277587890625],"size":[210,58.94380569458008],"flags":{},"order":5,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Define model"],"color":"#432","bgcolor":"#653"},{"id":80,"type":"PtnResnetModel","pos":[2478.35009765625,1413.751953125],"size":[379.3387145996094,113.2640609741211],"flags":{},"order":6,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[91,104],"slot_index":0}],"properties":{"Node name for S&R":"PtnResnetModel"},"widgets_values":["(3,32,32)",10,10]},{"id":63,"type":"PtDataLoader","pos":[2473.048828125,1723.90625],"size":[400,200],"flags":{},"order":12,"mode":0,"inputs":[{"name":"dataset","type":"PTVDATASET","link":74}],"outputs":[{"name":"PTDATALOADER","type":"PTDATALOADER","links":[105],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[128,true,"{\"num_workers\":1}"]},{"id":64,"type":"PtDataLoader","pos":[2473.048828125,2039.8226318359375],"size":[400,200],"flags":{},"order":14,"mode":0,"inputs":[{"name":"dataset","type":"PTVDATASET","link":75}],"outputs":[{"name":"PTDATALOADER","type":"PTDATALOADER","links":[106],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[128,false,"{\"num_workers\":1}"]},{"id":55,"type":"MarkdownNote","pos":[3030.089599609375,1947.71044921875],"size":[232.82992553710938,58],"flags":{},"order":7,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Train"],"color":"#432","bgcolor":"#653"},{"id":91,"type":"PtTrainClassificationModelLr","pos":[3023.411376953125,1655.6585693359375],"size":[403.1999816894531,234],"flags":{},"order":19,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":104},{"name":"train_loader","type":"PTDATALOADER","link":105},{"name":"optimizer","type":"PTOPTIMIZER","link":107},{"name":"scheduler","type":"PTLRSCHEDULER","link":112},{"name":"val_loader","type":"PTDATALOADER","link":106,"shape":7}],"outputs":[{"name":"Model","type":"PTMODEL","links":[108],"slot_index":0},{"name":"train loss","type":"TENSOR","links":null},{"name":"val loss","type":"TENSOR","links":null}],"properties":{"Node name for S&R":"PtTrainClassificationModelLr"},"widgets_values":[200,true,false,10,false]},{"id":11,"type":"PtoAdam","pos":[2976.371337890625,1387.1434326171875],"size":[246.16842651367188,110.40398406982422],"flags":{},"order":11,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":91}],"outputs":[{"name":"PTOPTIMIZER","type":"PTOPTIMIZER","links":[107,111],"slot_index":0}],"properties":{"Node name for S&R":"PtoAdam"},"widgets_values":[0.001,0.9,0.999]},{"id":93,"type":"PtoLrSchedulerCosineAnnealing","pos":[3274.407958984375,1388.4407958984375],"size":[280.896240234375,86.71690368652344],"flags":{},"order":16,"mode":0,"inputs":[{"name":"optimizer","type":"PTOPTIMIZER","link":111}],"outputs":[{"name":"PTLRSCHEDULER","type":"PTLRSCHEDULER","links":[112],"slot_index":0}],"properties":{"Node name for S&R":"PtoLrSchedulerCosineAnnealing"},"widgets_values":[200,0]},{"id":89,"type":"CDAStringCreate","pos":[3536.05712890625,1534.6680908203125],"size":[385.5651550292969,59.037696838378906],"flags":{},"order":8,"mode":0,"inputs":[],"outputs":[{"name":"STRING","type":"STRING","links":[99],"slot_index":0}],"properties":{"Node name for S&R":"CDAStringCreate"},"widgets_values":["/container/transfer/models/cifar10_v11_2nd_run.pt"]}],"links":[[70,19,0,51,0,"STRING"],[74,61,0,63,0,"PTVDATASET"],[75,62,0,64,0,"PTVDATASET"],[83,61,0,76,0,"PTVDATASET"],[84,76,0,77,0,"INT"],[85,78,0,79,0,"INT"],[86,62,0,78,0,"PTVDATASET"],[89,81,0,62,0,"PTVTRANSFORM"],[91,80,0,11,0,"PTMODEL"],[99,89,0,19,1,"STRING"],[103,90,0,61,0,"PTVTRANSFORM"],[104,80,0,91,0,"PTMODEL"],[105,63,0,91,1,"PTDATALOADER"],[106,64,0,91,4,"PTDATALOADER"],[107,11,0,91,2,"PTOPTIMIZER"],[108,91,0,19,0,"PTMODEL"],[111,11,0,93,0,"PTOPTIMIZER"],[112,93,0,91,3,"PTLRSCHEDULER"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.3405679149452571,"offset":[-65.83749226699729,-475.28970703078994]},"node_versions":{"pt-wrapper":"eb9d05ed533f01888758f204deea434badf4a0a9","data-analysis":"bde4a07755390554f858e0b8fd8554a3123a6612"},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/examples/workflows/dog_cat_classification_prediction.json b/examples/workflows/dog_cat_classification_prediction.json deleted file mode 100644 index 187c32a..0000000 --- a/examples/workflows/dog_cat_classification_prediction.json +++ /dev/null @@ -1 +0,0 @@ -{"last_node_id":74,"last_link_id":97,"nodes":[{"id":73,"type":"PtFromImageTranspose","pos":[1488.2567138671875,1628.712646484375],"size":[289.79998779296875,26],"flags":{},"order":4,"mode":0,"inputs":[{"name":"image","type":"IMAGE","link":95}],"outputs":[{"name":"TENSOR","type":"TENSOR","links":[96],"slot_index":0}],"properties":{"Node name for S&R":"PtFromImageTranspose"}},{"id":74,"type":"PtInterpolateToSize","pos":[1801.049560546875,1628.712646484375],"size":[315,106],"flags":{},"order":5,"mode":0,"inputs":[{"name":"tens","type":"TENSOR","link":96}],"outputs":[{"name":"TENSOR","type":"TENSOR","links":[97],"slot_index":0}],"properties":{"Node name for S&R":"PtInterpolateToSize"},"widgets_values":[256,256,"nearest-exact"]},{"id":69,"type":"PtSqueeze","pos":[2206.266357421875,1639.8197021484375],"size":[210,60.5162467956543],"flags":{},"order":6,"mode":0,"inputs":[{"name":"tens","type":"TENSOR","link":97}],"outputs":[{"name":"TENSOR","type":"TENSOR","links":[88],"slot_index":0}],"properties":{"Node name for S&R":"PtSqueeze"},"widgets_values":[0]},{"id":55,"type":"PtPredictClassificationModel","pos":[2482.278564453125,1626.3341064453125],"size":[405.5999755859375,98],"flags":{},"order":7,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":70},{"name":"inputs","type":"TENSOR","link":88},{"name":"class_id_to_name_map","type":"PYDICT","link":69}],"outputs":[{"name":"class_name","type":"STRING","links":[74],"slot_index":0},{"name":"class_id","type":"INT","links":null},{"name":"probability","type":"FLOAT","links":[93],"slot_index":2}],"properties":{"Node name for S&R":"PtPredictClassificationModel"},"widgets_values":[true]},{"id":56,"type":"PyStringToDict","pos":[2206.266357421875,1774.539306640625],"size":[210,81.30076599121094],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[{"name":"PYDICT","type":"PYDICT","links":[69],"slot_index":0}],"properties":{"Node name for S&R":"PyStringToDict"},"widgets_values":["{0:\"cat\", 1:\"dog\"}"]},{"id":60,"type":"PandasShowText","pos":[2970.7451171875,1629.65576171875],"size":[221.17684936523438,76],"flags":{},"order":8,"mode":0,"inputs":[{"name":"text","type":"STRING","link":74,"widget":{"name":"text"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"properties":{"Node name for S&R":"PandasShowText"},"widgets_values":["","cat"]},{"id":72,"type":"CDAShowFloat","pos":[2976.763916015625,1770.6387939453125],"size":[210,76],"flags":{},"order":9,"mode":0,"inputs":[{"name":"float_scalar","type":"FLOAT","link":93,"widget":{"name":"float_scalar"}}],"outputs":[{"name":"STRING","type":"STRING","links":null,"shape":6}],"properties":{"Node name for S&R":"CDAShowFloat"},"widgets_values":[0,"0.9118371605873108"]},{"id":23,"type":"PtLoadModel","pos":[2206.266357421875,1482.663330078125],"size":[210,75.45188903808594],"flags":{},"order":3,"mode":0,"inputs":[{"name":"model","type":"PTMODEL","link":62}],"outputs":[{"name":"model","type":"PTMODEL","links":[70],"slot_index":0}],"properties":{"Node name for S&R":"PtLoadModel"},"widgets_values":["dog_cat.pt"]},{"id":57,"type":"LoadImage","pos":[1135.1787109375,1628.712646484375],"size":[315,314],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[{"name":"IMAGE","type":"IMAGE","links":[95],"slot_index":0},{"name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["cat.jpg","image"]},{"id":48,"type":"PtnConvModel","pos":[1801.049560546875,1306.3406982421875],"size":[315,202],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","type":"PTMODEL","links":[62],"slot_index":0}],"properties":{"Node name for S&R":"PtnConvModel"},"widgets_values":["(3,256,256)",0,2,"[32,64,128,256,512]","[3,3,3,3,3]","[1,1,1,1,1]","[True,True,True,True,True]"]}],"links":[[62,48,0,23,0,"PTMODEL"],[69,56,0,55,2,"PYDICT"],[70,23,0,55,0,"PTMODEL"],[74,55,0,60,0,"STRING"],[88,69,0,55,1,"TENSOR"],[93,55,2,72,0,"FLOAT"],[95,57,0,73,0,"IMAGE"],[96,73,0,74,0,"TENSOR"],[97,74,0,69,0,"TENSOR"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.6339486400897979,"offset":[-879.509850098028,-1085.7559202042573]},"node_versions":{"data-analysis":"fd6fa94b931cd094ce5766852f5914d7291cb464","comfy-core":"v0.3.10-40-gd0f3752"},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/modules/pytorch_wrapper/pt_train_classification_model_lr.py b/modules/pytorch_wrapper/pt_train_classification_model_lr.py new file mode 100644 index 0000000..2dcf095 --- /dev/null +++ b/modules/pytorch_wrapper/pt_train_classification_model_lr.py @@ -0,0 +1,178 @@ +import os +import sys +from typing import Any, Dict, Optional, Tuple, Union +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np + +class PtTrainClassificationModelLr: + """ + Trains a classification model using a given dataset, optimizer, and number of epochs with learning rate decay. + + early_stopping_rounds: Specifies the number of rounds to monitor the evaluation loss when early_stopping is enabled and validation data is specified. If the loss does not decrease compared to the current best loss within this period, training will be stopped. + output_best_val_model: Whether to outputs the model with the best val loss or the model from the last epoch. + + 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", {}), + "scheduler": ("PTLRSCHEDULER", {}), + "epochs": ("INT", {"default":1, "min":1, "max":1e6}), + "use_gpu": ("BOOLEAN", {"default":False}), + "early_stopping": ("BOOLEAN", {"default":False}), + "early_stopping_rounds": ("INT", {"default":10, "min":1, "max":1000}), + "output_best_val_model": ("BOOLEAN", {"default":True}) + }, + "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, + scheduler: Union[torch.optim.lr_scheduler._LRScheduler, torch.optim.lr_scheduler.ReduceLROnPlateau], + epochs: int, + use_gpu: bool, + early_stopping: bool, + early_stopping_rounds: int, + output_best_val_model: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. + scheduler (Union[torch.optim.lr_scheduler._LRScheduler, torch.optim.lr_scheduler.ReduceLROnPlateau]): Learning rate scheduler. + epochs (int): Number of training epochs. + use_gpu (bool): Whether to use GPU for training. + early_stopping (bool): Whether to turn on early stopping + early_stopping_rounds (int): Specifies the number of rounds to monitor the evaluation loss when early_stopping is enabled and validation data is specified. If the loss does not decrease compared to the current best loss within this period, training will be stopped. + output_best_val_model (bool): Whether to outputs the model with the best val loss or the model from the last epoch. + 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. + """ + if (isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau) and + val_loader is None): + raise ValueError("Validation data loader needs to be specified when ReduceLROnPlateau scheduler is specified.") + + checkpoint_dir = os.path.realpath( + os.path.join(os.path.dirname(__file__), "..", "..", "models", "tmp")) + print(f"Checkpoint tmp directory is set to {checkpoint_dir}") + if os.path.exists(checkpoint_dir) is False: + os.makedirs(checkpoint_dir) + print(f"Created {checkpoint_dir}") + best_checkpoint_path = os.path.join(checkpoint_dir, "best_model.pt") + + best_val_loss = None + rounds_since_best_val_loss = 0 + + with torch.inference_mode(False): + train_loss = [] + val_loss = [] + if use_gpu: + model.to("cuda") + + loss_function = nn.CrossEntropyLoss() # mean loss + + # Train the model + for epoch in range(epochs): + model.train() + total_loss = 0 + total_samples = len(train_loader.dataset) + 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() + + batch_size = labels.size(0) + batch_loss = loss.item() * batch_size + total_loss += batch_loss + + epoch_loss = total_loss / total_samples + + print(f"Epoch (train) {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}") + train_loss.append(epoch_loss) + + if val_loader: + with torch.no_grad(): + model.train(False) # TODO: Change to model.*e*v*a*l() once Comfy's security checker is fixed. + total_loss = 0 + total_samples = len(val_loader.dataset) + 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) + batch_size = labels.size(0) + batch_loss = loss.item() * batch_size + total_loss += batch_loss + + epoch_loss = total_loss / total_samples + print(f"Epoch (val) {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}") + val_loss.append(epoch_loss) + + if best_val_loss is None or epoch_loss < best_val_loss: + best_val_loss = epoch_loss + rounds_since_best_val_loss = 0 + torch.save(model.state_dict(), best_checkpoint_path) + print(f"Updated best val loss. Saved the model as {best_checkpoint_path}") + else: + rounds_since_best_val_loss += 1 + if early_stopping: + if rounds_since_best_val_loss >= early_stopping_rounds: + print(f"Val loss did not improve in {early_stopping_rounds} rounds. Stopping training.") + break + + if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau): + scheduler.step(epoch_loss) + else: + scheduler.step() + + if use_gpu: + model.to("cpu") + + if val_loader: + if epoch_loss > best_val_loss and output_best_val_model: + sd = torch.load(best_checkpoint_path, weights_only=True) + model.load_state_dict(sd) + print("Loaded the best model for output") + else: + print("Using the final epoch model for output.") + + return (model, + torch.tensor(train_loss, dtype=torch.float32), + torch.tensor(val_loss, dtype=torch.float32)) \ No newline at end of file diff --git a/modules/pytorch_wrapper/pto_lr_scheduler_cosine_annealing.py b/modules/pytorch_wrapper/pto_lr_scheduler_cosine_annealing.py new file mode 100644 index 0000000..c3c58ed --- /dev/null +++ b/modules/pytorch_wrapper/pto_lr_scheduler_cosine_annealing.py @@ -0,0 +1,79 @@ +from typing import Any, Dict, Tuple +import torch +import torch.optim +from torch.optim.lr_scheduler import CosineAnnealingLR + + +class PtoLrSchedulerCosineAnnealing: + """ + Creates a cosine annealing learning rate scheduler for an optimizer. + + https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.CosineAnnealingLR.html + lists: + + eta_t = eta_min + (1/2) * (eta_max - eta_min) * (1 + cos((t_cur / t_max) * pi)) + + eta_min is set to 0 by default, so the equation simplifies to: + + eta_t = (1/2) * eta_max * (1 + cos((t_cur / t_max) * pi)) + + When t_cur = 0: + (1 + cos((t_cur / t_max) * pi)) evaluates to (1 + cos(0 * pi)) = (1 + 1) = 2 + + When t_cur = t_max: + (1 + cos((t_cur / t_max) * pi)) evaluates to (1 + cos(1 * pi)) = (1 - 1) = 0 + + Therefore, eta_t value changes as follows: + * Starts from (1/2) * eta_max * 2, which simplifies to eta_max + * Ends at 0 + + Parameters + + **optimizer** + (torch.optim.Optimizer) – Wrapped optimizer whose learning rate needs scheduling. + + **num_epochs** + (int) – Number of epochs. This maps to the T_max parameter for CosineAnnealingLR internally. + + **minimum_lr** + (float) – Minimum learning rate. + + 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": { + "optimizer": ("PTOPTIMIZER", {}), + "num_epochs": ("INT", {"default": 10, "min": 1, "max": int(1e8)}), + "minimum_lr": ("FLOAT", {"default": 0, "min": 0, "max": 1.0, "step": 1e-8}), + } + } + + RETURN_TYPES: Tuple[str, ...] = ("PTLRSCHEDULER",) + FUNCTION: str = "f" + CATEGORY: str = "Training" + + def f(self, optimizer: torch.optim.Optimizer, num_epochs: int, minimum_lr: float) -> Tuple[CosineAnnealingLR]: + """ + Creates a cosine annealing learning rate scheduler for an optimizer. + + Args: + optimizer (torch.optim.Optimizer): The optimizer for which to schedule the learning rate. + num_epochs (int): Number of epochs. + gamma (float): Minimum learning rate. + + Returns: + tuple: A tuple containing the CosineAnnealingLR scheduler. + """ + with torch.inference_mode(False): + lrs = CosineAnnealingLR(optimizer, T_max=num_epochs, eta_min=minimum_lr) + return (lrs,) diff --git a/modules/pytorch_wrapper/pto_lr_scheduler_reduce_on_plateau.py b/modules/pytorch_wrapper/pto_lr_scheduler_reduce_on_plateau.py new file mode 100644 index 0000000..6e13729 --- /dev/null +++ b/modules/pytorch_wrapper/pto_lr_scheduler_reduce_on_plateau.py @@ -0,0 +1,60 @@ +from typing import Any, Dict, Tuple +import torch +import torch.optim +from torch.optim.lr_scheduler import ReduceLROnPlateau + + +class PtoLrSchedulerReduceOnPlateau: + """ + Creates a reduce-on-plateau learning rate scheduler for an optimizer. + + Parameters + + **optimizer** + (torch.optim.Optimizer) – Wrapped optimizer whose learning rate needs scheduling. + + **grace_period** + (int) – Number of epochs to monitor the val loss reduction. If loss does not decrease within grace period epochs, learning rate will be reduced by multiplying the gamma. E.g. if gamma is set to 0.5, new learning rate will be 0.5 * current learning rate. + + **gamma** + (float) – Learning rate multiplier to be applied when grace period is up. + + 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": { + "optimizer": ("PTOPTIMIZER", {}), + "grace_period": ("INT", {"default": 10, "min": 0, "max": int(1e8)}), + "gamma": ("FLOAT", {"default": 0.1, "min": 1e-6, "max": 1.0, "step": 1e-8}) + } + } + + RETURN_TYPES: Tuple[str, ...] = ("PTLRSCHEDULER",) + FUNCTION: str = "f" + CATEGORY: str = "Training" + + def f(self, optimizer: torch.optim.Optimizer, grace_period: int, gamma: float) -> Tuple[ReduceLROnPlateau]: + """ + Creates a reduce-on-plateau learning rate scheduler for an optimizer. + + Args: + optimizer (torch.optim.Optimizer): The optimizer for which to schedule the learning rate. + grace_period (int): Number of epochs to monitor for val loss reduction. + gamma (float): Multiplicative factor of learning rate decay. + + Returns: + tuple: A tuple containing the CosineAnnealingLR scheduler. + """ + with torch.inference_mode(False): + lrs = ReduceLROnPlateau(optimizer, mode="min", factor=gamma, patience=grace_period) + return (lrs,) diff --git a/modules/pytorch_wrapper/pto_lr_scheduler_step.py b/modules/pytorch_wrapper/pto_lr_scheduler_step.py new file mode 100644 index 0000000..ae0c056 --- /dev/null +++ b/modules/pytorch_wrapper/pto_lr_scheduler_step.py @@ -0,0 +1,58 @@ +from typing import Any, Dict, Tuple +import torch +import torch.optim +from torch.optim.lr_scheduler import StepLR + +class PtoLrSchedulerStep: + """ + Creates a StepLR learning rate scheduler for an optimizer. + + Parameters + + **optimizer** + (torch.optim.Optimizer) – Wrapped optimizer whose learning rate needs scheduling. + + **step_size** + (int) – Period of learning rate decay. + + **gamma** + (float) – Multiplicative factor of learning rate decay. + + 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": { + "optimizer": ("PTOPTIMIZER", {}), + "step_size": ("INT", {"default": 10, "min": 1, "max": int(1e8)}), + "gamma": ("FLOAT", {"default": 0.1, "min": 1e-6, "max": 1.0, "step": 1e-8}), + } + } + + RETURN_TYPES: Tuple[str, ...] = ("PTLRSCHEDULER",) + FUNCTION: str = "f" + CATEGORY: str = "Training" + + def f(self, optimizer: torch.optim.Optimizer, step_size: int, gamma: float) -> Tuple[StepLR]: + """ + Creates a StepLR learning rate scheduler. + + Args: + optimizer (torch.optim.Optimizer): The optimizer for which to schedule the learning rate. + step_size (int): Period of learning rate decay. + gamma (float): Multiplicative factor of learning rate decay. + + Returns: + tuple: A tuple containing the StepLR scheduler. + """ + with torch.inference_mode(False): + lrs = StepLR(optimizer, step_size=step_size, gamma=gamma) + return (lrs,) diff --git a/pyproject.toml b/pyproject.toml index 3796f4b..5f21442 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "pt-wrapper" description = "PyTorch extension for ComfyUI featuring extensive PyTorch wrapper nodes for seamless tensor operations and PyTorch model training." -version = "3.0.5" +version = "3.0.6" license = {file = "LICENSE"} [project.urls] diff --git a/tests/pytorch_wrapper/pt_train_classification_conv_model_lr_test.py b/tests/pytorch_wrapper/pt_train_classification_conv_model_lr_test.py new file mode 100644 index 0000000..28f6012 --- /dev/null +++ b/tests/pytorch_wrapper/pt_train_classification_conv_model_lr_test.py @@ -0,0 +1,77 @@ +import os +import sys +import unittest +import torch +import torch.nn as nn + +PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", "..")) +MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules") +sys.path.append(MODULE_ROOT) + +from pytorch_wrapper.ptn_conv_model import PtnConvModel +from pytorch_wrapper.pto_adam import PtoAdam +from pytorch_wrapper.ptv_dataset import PtvDataset +from pytorch_wrapper.ptv_transforms_to_tensor import PtvTransformsToTensor +from pytorch_wrapper.pt_data_loader import PtDataLoader +from pytorch_wrapper.pt_train_classification_model_lr import PtTrainClassificationModelLr +from pytorch_wrapper.pt_save_model import PtSaveModel +from pytorch_wrapper.pto_lr_scheduler_step import PtoLrSchedulerStep # Updated import + + +class TestPtTrainClassificationConvModelLr(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + compose_node = PtvTransformsToTensor() + transform = compose_node.f()[0] + dataset_node = PtvDataset() + dataset = dataset_node.f("FashionMNIST", + True, # download + "", # root + '{"train": True}', + transform=transform)[0] + data_loader_node = PtDataLoader() + self.train_loader = data_loader_node.f( + dataset=dataset, + batch_size=32, + shuffle=True, + parameters='{"num_workers":1}' + )[0] + + self.model_node = PtnConvModel() + self.model = self.model_node.f( + input_dim="(1, 28, 28)", + penultimate_dim=0, + output_dim=10, + channel_list="[32,64,128,256]", + kernel_size_list="[3,3,3,3]", + padding_list="[1,1,1,1]", + downsample_list="[True,True,True,True]" + )[0] + + self.optimizer = PtoAdam().f(self.model, 0.001, 0.9, 0.999)[0] + + self.scheduler = PtoLrSchedulerStep().f(self.optimizer, + 3, # steps + 0.1 # gamma + )[0] + + self.trainer = PtTrainClassificationModelLr() + self.save_model = PtSaveModel() + + def test_1(self): + epochs = 4 + trained_model = self.trainer.f( + self.model, + self.train_loader, + self.optimizer, + self.scheduler, + epochs, # epochs + use_gpu=True, + early_stopping=False, + early_stopping_rounds=1, + output_best_val_model=False)[0] + self.save_model.f(trained_model, f"fashion_mnist_{epochs}_epochs_conv.pt") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/pto_lr_scheduler_cosine_annealing_test.py b/tests/pytorch_wrapper/pto_lr_scheduler_cosine_annealing_test.py new file mode 100644 index 0000000..362bdd0 --- /dev/null +++ b/tests/pytorch_wrapper/pto_lr_scheduler_cosine_annealing_test.py @@ -0,0 +1,41 @@ +import os +import sys +import unittest +import torch +import torch.optim as optim +import math + +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.pto_lr_scheduler_cosine_annealing import PtoLrSchedulerCosineAnnealing + + +class TestPtoLrSchedulerCosineAnnealing(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.node = PtoLrSchedulerCosineAnnealing() + + # Create a pseudo optimizer with learning rate parameter + self.model = torch.nn.Linear(2, 2) # Dummy model + self.optimizer = optim.SGD(self.model.parameters(), lr=0.1) + self.num_epochs = 10 + self.minimum_lr = 0.01 + + def test_lr_scheduler_cosine_annealing(self): + """Test if CosineAnnealingLR correctly updates learning rate.""" + scheduler, = self.node.f(self.optimizer, self.num_epochs, self.minimum_lr) + + # Ensure returned object is CosineAnnealingLR + self.assertIsInstance(scheduler, optim.lr_scheduler.CosineAnnealingLR, "Returned object is not a CosineAnnealingLR scheduler") + + initial_lr = self.optimizer.param_groups[0]['lr'] + for t in range(1, self.num_epochs + 1): + scheduler.step() + updated_lr = self.optimizer.param_groups[0]['lr'] + expected_lr = self.minimum_lr + 0.5 * (initial_lr - self.minimum_lr) * (1 + math.cos(math.pi * t / self.num_epochs)) + self.assertAlmostEqual(updated_lr, expected_lr, places=6, msg=f"Expected LR {expected_lr}, got {updated_lr} at step {t}") + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/pto_lr_scheduler_reduce_on_plateau_test.py b/tests/pytorch_wrapper/pto_lr_scheduler_reduce_on_plateau_test.py new file mode 100644 index 0000000..2a4ed34 --- /dev/null +++ b/tests/pytorch_wrapper/pto_lr_scheduler_reduce_on_plateau_test.py @@ -0,0 +1,65 @@ +import os +import sys +import unittest +import torch +import torch.optim as optim + +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.pto_lr_scheduler_reduce_on_plateau import PtoLrSchedulerReduceOnPlateau # Updated import + + +class TestPtoLrSchedulerReduceOnPlateau(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.node = PtoLrSchedulerReduceOnPlateau() + + # Create a pseudo optimizer with learning rate parameter + self.model = torch.nn.Linear(2, 2) # Dummy model + self.optimizer = optim.SGD(self.model.parameters(), lr=0.1) + self.grace_period = 3 + self.gamma = 0.5 + + def test_lr_scheduler_reduce_on_plateau(self): + """Test if ReduceLROnPlateau correctly updates learning rate.""" + scheduler, = self.node.f(self.optimizer, self.grace_period, self.gamma) + + # Ensure returned object is ReduceLROnPlateau + self.assertIsInstance(scheduler, optim.lr_scheduler.ReduceLROnPlateau, "Returned object is not a ReduceLROnPlateau scheduler") + + initial_lr = self.optimizer.param_groups[0]['lr'] + # No improvement in loss + # Grace period map: 0, 1 2 3 + # First epoch is the baseline, so it does not count. + # In this test case, the number 4 (the 4th) bad loss is greater than grace_period=3, + # so the loss reduction happens there. + loss_values = [0.5, 0.5, 0.5, 0.5, 0.5] + for loss in loss_values: + scheduler.step(loss) + updated_lr = self.optimizer.param_groups[0]['lr'] + + # Check if learning rate was updated correctly after grace period + expected_lr = initial_lr * self.gamma + self.assertAlmostEqual(updated_lr, expected_lr, places=6, msg=f"Expected LR {expected_lr}, got {updated_lr}") + + def test_lr_scheduler_reduce_on_plateau_no_apply(self): + """Test if ReduceLROnPlateau correctly does not prematurely update learning rate.""" + scheduler, = self.node.f(self.optimizer, self.grace_period, self.gamma) + + # Ensure returned object is ReduceLROnPlateau + self.assertIsInstance(scheduler, optim.lr_scheduler.ReduceLROnPlateau, "Returned object is not a ReduceLROnPlateau scheduler") + + initial_lr = self.optimizer.param_groups[0]['lr'] + loss_values = [0.5, 0.5, 0.5, 0.5] + for loss in loss_values: + scheduler.step(loss) + updated_lr = self.optimizer.param_groups[0]['lr'] + + expected_lr = initial_lr + self.assertAlmostEqual(updated_lr, expected_lr, places=6, msg=f"Expected LR {expected_lr}, got {updated_lr}") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/pto_lr_scheduler_step_test.py b/tests/pytorch_wrapper/pto_lr_scheduler_step_test.py new file mode 100644 index 0000000..56ee104 --- /dev/null +++ b/tests/pytorch_wrapper/pto_lr_scheduler_step_test.py @@ -0,0 +1,44 @@ +import os +import sys +import unittest +import torch +import torch.optim as optim + +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.pto_lr_scheduler_step import PtoLrSchedulerStep # Updated import + + +class TestPtoLrSchedulerStep(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.node = PtoLrSchedulerStep() + + # Create a pseudo optimizer with learning rate parameter + # For every 5 steps, if gamma is 0.1, then lr *= 0.5 + # so 0.1 -> 0.05 -> 0.025 + self.model = torch.nn.Linear(2, 2) # Dummy model + self.optimizer = optim.SGD(self.model.parameters(), lr=0.1) + self.step_size = 5 + self.gamma = 0.1 + + def test_lr_scheduler_step(self): + """Test if learning rate scheduler correctly updates learning rate.""" + scheduler, = self.node.f(self.optimizer, self.step_size, self.gamma) + + # Ensure returned object is StepLR + self.assertIsInstance(scheduler, optim.lr_scheduler.StepLR, "Returned object is not a StepLR scheduler") + + initial_lr = self.optimizer.param_groups[0]['lr'] + for _ in range(self.step_size): + scheduler.step() + updated_lr = self.optimizer.param_groups[0]['lr'] + + # Check if learning rate was updated correctly + expected_lr = initial_lr * self.gamma + self.assertAlmostEqual(updated_lr, expected_lr, places=6, msg=f"Expected LR {expected_lr}, got {updated_lr}") + +if __name__ == "__main__": + unittest.main()