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

Added lr scheduler nodes and cifar10 workflows.

Browse Source
This commit is contained in:
HowToSD
2025-03-12 14:57:44 -07:00
parent 635c97c037
commit acb5d9f820
20 changed files with 770 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1 -1
View File
@@ -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)
__init__.py
+12
View File
@@ -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",
docs/model_training.md
+8 -3
View File
@@ -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.**
docs/reference/node_reference.md
+4
View File
@@ -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. |
docs/reference/pt_train_classification_model_lr.md
+31
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Training
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pto_lr_scheduler_cosine_annealing.md
+49
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Training
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pto_lr_scheduler_reduce_on_plateau.md
+30
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Training
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pto_lr_scheduler_step.md
+30
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Training
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
examples/workflows/cifar10_eval_v11_2nd_run.json
+1
View File
File diff suppressed because one or more lines are too long
examples/workflows/cifar10_train_v11_2nd_run.json
+1
View File
File diff suppressed because one or more lines are too long
examples/workflows/dog_cat_classification_prediction.json
-1
View File
@@ -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}
modules/pytorch_wrapper/pt_train_classification_model_lr.py
+178
View File
@@ -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))
modules/pytorch_wrapper/pto_lr_scheduler_cosine_annealing.py
+79
View File
@@ -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,)
modules/pytorch_wrapper/pto_lr_scheduler_reduce_on_plateau.py
+60
View File
@@ -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,)
modules/pytorch_wrapper/pto_lr_scheduler_step.py
+58
View File
@@ -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,)
pyproject.toml
+1 -1
View File
@@ -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]
tests/pytorch_wrapper/pt_train_classification_conv_model_lr_test.py
+77
View File
@@ -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()
tests/pytorch_wrapper/pto_lr_scheduler_cosine_annealing_test.py
+41
View File
@@ -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()
tests/pytorch_wrapper/pto_lr_scheduler_reduce_on_plateau_test.py
+65
View File
@@ -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()
tests/pytorch_wrapper/pto_lr_scheduler_step_test.py
+44
View File
@@ -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()