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

Added Pto AdamW and Pto SGD optimizer nodes.

Browse Source
This commit is contained in:
HowToSD
2025-03-15 15:59:40 -07:00
parent a98448c176
commit 7197708ee9
14 changed files with 331 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+6
View File
@@ -127,9 +127,11 @@ from .modules.pytorch_wrapper.ptn_linear import PtnLinear
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_adamw import PtoAdamW
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.pto_sgd import PtoSGD
from .modules.pytorch_wrapper.pto_simple import PtoSimple
from .modules.pytorch_wrapper.ptv_dataset import PtvDataset
from .modules.pytorch_wrapper.ptv_dataset_len import PtvDatasetLen
@@ -274,9 +276,11 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
"PtnLinearModel": PtnLinearModel,
"PtnResnetModel": PtnResnetModel,
"PtoAdam": PtoAdam,
"PtoAdamW": PtoAdamW,
"PtoLrSchedulerCosineAnnealing": PtoLrSchedulerCosineAnnealing,
"PtoLrSchedulerReduceOnPlateau": PtoLrSchedulerReduceOnPlateau,
"PtoLrSchedulerStep": PtoLrSchedulerStep,
"PtoSGD": PtoSGD,
"PtoSimple": PtoSimple,
"PtvDataset": PtvDataset,
"PtvDatasetLen": PtvDatasetLen,
@@ -421,9 +425,11 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
"PtnLinearModel": "Ptn Linear Model",
"PtnResnetModel": "Ptn Resnet Model",
"PtoAdam": "Pto Adam",
"PtoAdamW": "Pto AdamW",
"PtoLrSchedulerCosineAnnealing": "Pto Lr Scheduler Cosine Annealing",
"PtoLrSchedulerReduceOnPlateau": "Pto Lr Scheduler Reduce On Plateau",
"PtoLrSchedulerStep": "Pto Lr Scheduler Step",
"PtoSGD": "Pto SGD",
"PtoSimple": "Pto Simple",
"PtvDataset": "Ptv Dataset",
"PtvDatasetLen": "Ptv Dataset Len",
dev/scripts/init_py_generate.py
+2
View File
@@ -78,9 +78,11 @@ def format_display_name(class_name):
("MP L", "MPL"),
("SN S", "SNS"),
(" N A", " NA"),
("Adam W", "AdamW"),
("Re L U", "ReLU"),
("Si L U", "SiLU"),
("G E L U", "GELU"),
("S G D", "SGD"),
("Data Frame", "DataFrame"),
("Multi Index", "MultiIndex"),
("MPL Bar", "MPL Bar Chart"),
docs/reference/node_reference.md
+7 -2
View File
@@ -110,6 +110,13 @@
| [Pt Einsum](pt_einsum.md) | Performs Tensor operations specified in the Einstein summation equation. |
| [Pt Mat Mul](pt_matmul.md) | Performs matrix multiplication on two PyTorch tensors. |
| [Pt Mm](pt_mm.md) | Performs 2D matrix multiplication on two PyTorch tensors. |
## PyTorch wrapper - Optimizer
| Node | Description |
| --- | --- |
| [Pto Adam](pto_adam.md) | Instantiates the Adam optimizer. |
| [Pto AdamW](pto_adamw.md) | Instantiates the AdamW optimizer. |
| [Pto SGD](pto_sgd.md) | Instantiates the SGD optimizer. |
| [Pto Simple](pto_simple.md) | Instantiates the most basic optimizer to update W using the below formula: |
## PyTorch wrapper - Reduction operation & Summary statistics
| Node | Description |
| --- | --- |
@@ -170,11 +177,9 @@
| [Ptn Linear](ptn_linear.md) | A linear model consisting of a single dense layer. |
| [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. |
| [Pto Simple](pto_simple.md) | Instantiates the most basic optimizer to update W using the below formula: |
| [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/pto_adam.md
+1 -1
View File
@@ -15,6 +15,6 @@ Instantiates the Adam optimizer.
| Ptoptimizer |
<HR>
Category: PyTorch wrapper - Training
Category: PyTorch wrapper - Optimizer
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pto_adamw.md
+22
View File
@@ -0,0 +1,22 @@
# Pto AdamW
Instantiates the AdamW optimizer.
## Input
| Name | Data type |
|---|---|
| model | Ptmodel |
| learning_rate | Float |
| beta1 | Float |
| beta2 | Float |
| weight_decay | Float |
| amsgrad | Boolean |
## Output
| Data type |
|---|
| Ptoptimizer |
<HR>
Category: PyTorch wrapper - Optimizer
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pto_sgd.md
+30
View File
@@ -0,0 +1,30 @@
# Pto SGD
Instantiates the SGD optimizer.
Parameters
model (PTMODEL): The PyTorch model whose parameters will be optimized.
learning_rate (float): The learning rate for the AdamW optimizer.
momentum (float): Coefficient to apply to the past gradient to adjust the contribution of past gradients
dampening (float): Coefficient to adjust the contribution of current gradient.
weight_decay (float): The weight decay parameter.
nesterov (bool): If True, uses the Nesterov version.
## Input
| Name | Data type |
|---|---|
| model | Ptmodel |
| learning_rate | Float |
| momentum | Float |
| dampening | Float |
| weight_decay | Float |
| nesterov | Boolean |
## Output
| Data type |
|---|
| Ptoptimizer |
<HR>
Category: PyTorch wrapper - Optimizer
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pto_simple.md
+3 -2
View File
@@ -1,5 +1,6 @@
# Pto Simple
Instantiates the most basic optimizer to update W using the below formula:
Instantiates the most basic optimizer to update W using the below formula:
$$
W -= lr * G
$$
@@ -16,6 +17,6 @@ $$
| Ptoptimizer |
<HR>
Category: PyTorch wrapper - Training
Category: PyTorch wrapper - Optimizer
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
modules/pytorch_wrapper/pto_adam.py
+6 -4
View File
@@ -7,7 +7,7 @@ class PtoAdam:
"""
Instantiates the Adam optimizer.
category: PyTorch wrapper - Training
category: PyTorch wrapper - Optimizer
"""
@classmethod
@@ -36,10 +36,12 @@ class PtoAdam:
Instantiates the optimizer.
Args:
model (torch.nn.Module): PyTorch model
model (torch.nn.Module): The PyTorch model whose parameters will be optimized.
learning_rate (float): The learning rate for the AdamW optimizer.
beta1 (float): Exponential decay rate for the first moment estimates (moving average of gradients).
beta2 (float): Exponential decay rate for the second moment estimates (moving average of squared gradients).
Returns:
tuple: A tuple containing the ndarray.
tuple: A tuple containing the optimizer.
"""
with torch.inference_mode(False):
opt = torch.optim.Adam(model.parameters(), lr=learning_rate, betas=(beta1, beta2))
modules/pytorch_wrapper/pto_adamw.py
+62
View File
@@ -0,0 +1,62 @@
from typing import Any, Dict
import torch
import torch.nn as nn
import torch.nn.functional as F
class PtoAdamW:
"""
Pto AdamW:
Instantiates the AdamW optimizer.
category: PyTorch wrapper - Optimizer
"""
@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}),
"weight_decay": ("FLOAT", {"default":0.01, "min":0.0, "max":1.0, "step":0.0000000001}),
"amsgrad": ("BOOLEAN", {"default":False})
}
}
RETURN_TYPES: tuple = ("PTOPTIMIZER",)
FUNCTION: str = "f"
CATEGORY: str = "Training"
def f(self,
model:torch.nn.Module,
learning_rate:float,
beta1:float,
beta2:float,
weight_decay:float,
amsgrad: bool) -> tuple:
"""
Instantiates the optimizer.
Args:
model (torch.nn.Module): The PyTorch model whose parameters will be optimized.
learning_rate (float): The learning rate for the AdamW optimizer.
beta1 (float): Exponential decay rate for the first moment estimates (moving average of gradients).
beta2 (float): Exponential decay rate for the second moment estimates (moving average of squared gradients).
weight_decay (float): The weight decay parameter.
amsgrad (bool): If True, uses the AMSGrad version.
Returns:
tuple: A tuple containing the optimizer.
"""
with torch.inference_mode(False):
opt = torch.optim.AdamW(model.parameters(),
lr=learning_rate, betas=(beta1, beta2),
weight_decay=weight_decay,
amsgrad=amsgrad)
return(opt,)
modules/pytorch_wrapper/pto_sgd.py
+74
View File
@@ -0,0 +1,74 @@
from typing import Any, Dict
import torch
import torch.nn as nn
import torch.nn.functional as F
class PtoSGD:
"""
Pto SGD:
Instantiates the SGD optimizer.
Parameters
model (PTMODEL): The PyTorch model whose parameters will be optimized.
learning_rate (float): The learning rate for the AdamW optimizer.
momentum (float): Coefficient to apply to the past gradient to adjust the contribution of past gradients
dampening (float): Coefficient to adjust the contribution of current gradient.
weight_decay (float): The weight decay parameter.
nesterov (bool): If True, uses the Nesterov version.
category: PyTorch wrapper - Optimizer
"""
@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}),
"momentum": ("FLOAT", {"default":0.0, "min":0.0, "max":1.0, "step":0.0000000001}),
"dampening": ("FLOAT", {"default":0.0, "min":0.0, "max":1.0, "step":0.0000000001}),
"weight_decay": ("FLOAT", {"default":0.0, "min":0.0, "max":1.0, "step":0.0000000001}),
"nesterov": ("BOOLEAN", {"default":False})
}
}
RETURN_TYPES: tuple = ("PTOPTIMIZER",)
FUNCTION: str = "f"
CATEGORY: str = "Training"
def f(self,
model:torch.nn.Module,
learning_rate:float,
momentum:float,
dampening:float,
weight_decay:float,
nesterov: bool) -> tuple:
"""
Instantiates the optimizer.
Refer to https://pytorch.org/docs/stable/generated/torch.optim.SGD.html for further details.
Args:
model (torch.nn.Module): The PyTorch model whose parameters will be optimized.
learning_rate (float): The learning rate for the AdamW optimizer.
momentum (float): Coefficient to apply to the past gradient to adjust the contribution of past gradients
dampening (float): Coefficient to adjust the contribution of current gradient.
weight_decay (float): The weight decay parameter.
nesterov (bool): If True, uses the Nesterov version.
Returns:
tuple: A tuple containing the optimizer.
"""
with torch.inference_mode(False):
opt = torch.optim.SGD(model.parameters(),
lr=learning_rate,
momentum=momentum,
dampening=dampening,
weight_decay=weight_decay,
nesterov=nesterov)
return(opt,)
modules/pytorch_wrapper/pto_simple.py
+7 -3
View File
@@ -5,7 +5,9 @@ import torch.optim
class SimpleOptimizer(torch.optim.Optimizer):
"""
Implementation of the most basic optimizer to update W using the below formula:
Implementation of the most basic optimizer.
This optimizer updates W using the below formula:
$$
W -= lr * G
@@ -47,11 +49,13 @@ $$
class PtoSimple:
"""
Pto Simple:
Instantiates the most basic optimizer to update W using the below formula:
Instantiates the most basic optimizer to update W using the below formula:
$$
W -= lr * G
$$
category: PyTorch wrapper - Training
category: PyTorch wrapper - Optimizer
"""
@classmethod
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.10"
version = "3.0.11"
license = {file = "LICENSE"}
[project.urls]
tests/pytorch_wrapper/pto_adamw_test.py
+54
View File
@@ -0,0 +1,54 @@
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.pto_adamw import PtoAdamW
class TestPtoAdamW(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtoAdamW()
def test_adamw(self):
model = torch.nn.Linear(2, 3, bias=False)
torch.nn.init.constant_(model.weight, 5.0)
model.weight.grad = torch.full_like(model.weight, 1.0)
# AdamW hyperparameters
learning_rate = 0.1
beta1 = 0.9
beta2 = 0.99
weight_decay = 0.01
eps = 1e-8 # Default AdamW epsilon
opt = self.node.f(model, learning_rate=learning_rate, beta1=beta1, beta2=beta2,
weight_decay=weight_decay, amsgrad=False)[0]
expected_weight = torch.tensor(
[
[4.8950, 4.8950],
[4.8950, 4.8950],
[4.8950, 4.8950]
]
)
opt.step()
opt.zero_grad()
result = model.weight.data
self.assertTrue(
torch.allclose(result, expected_weight, atol=1e-6),
msg=f"Expected {expected_weight}, got {result}"
)
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pto_sgd_test.py
+56
View File
@@ -0,0 +1,56 @@
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.pto_sgd import PtoSGD
class TestPtoSGD(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtoSGD()
def test_sgd(self):
model = torch.nn.Linear(2, 3, bias=False)
torch.nn.init.constant_(model.weight, 5.0)
model.weight.grad = torch.full_like(model.weight, 1.0)
# Hyperparameters
learning_rate = 0.1
mu = 0.9
tau = 0.9
weight_decay = 0.01
opt = self.node.f(model, learning_rate=learning_rate,
momentum=mu,
dampening=tau,
weight_decay=weight_decay,
nesterov=False)[0]
expected_weight = torch.tensor(
[
[4.8950, 4.8950],
[4.8950, 4.8950],
[4.8950, 4.8950]
]
)
opt.step()
opt.zero_grad()
result = model.weight.data
self.assertTrue(
torch.allclose(result, expected_weight, atol=1e-6),
msg=f"Expected {expected_weight}, got {result}"
)
if __name__ == "__main__":
unittest.main()