diff --git a/__init__.py b/__init__.py index 085535f..fb852d5 100644 --- a/__init__.py +++ b/__init__.py @@ -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", diff --git a/dev/scripts/init_py_generate.py b/dev/scripts/init_py_generate.py index c33315c..32c511b 100644 --- a/dev/scripts/init_py_generate.py +++ b/dev/scripts/init_py_generate.py @@ -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"), diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index a85367a..232a4a0 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -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. | diff --git a/docs/reference/pto_adam.md b/docs/reference/pto_adam.md index eab0510..6b76c2b 100644 --- a/docs/reference/pto_adam.md +++ b/docs/reference/pto_adam.md @@ -15,6 +15,6 @@ Instantiates the Adam optimizer. | Ptoptimizer |
-Category: PyTorch wrapper - Training +Category: PyTorch wrapper - Optimizer ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/pto_adamw.md b/docs/reference/pto_adamw.md new file mode 100644 index 0000000..bfb1bd8 --- /dev/null +++ b/docs/reference/pto_adamw.md @@ -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 | + +
+Category: PyTorch wrapper - Optimizer + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/pto_sgd.md b/docs/reference/pto_sgd.md new file mode 100644 index 0000000..4ff22d9 --- /dev/null +++ b/docs/reference/pto_sgd.md @@ -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 | + +
+Category: PyTorch wrapper - Optimizer + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/pto_simple.md b/docs/reference/pto_simple.md index 71852d1..25768a2 100644 --- a/docs/reference/pto_simple.md +++ b/docs/reference/pto_simple.md @@ -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 |
-Category: PyTorch wrapper - Training +Category: PyTorch wrapper - Optimizer ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/modules/pytorch_wrapper/pto_adam.py b/modules/pytorch_wrapper/pto_adam.py index 5b8a9cf..8cbf101 100644 --- a/modules/pytorch_wrapper/pto_adam.py +++ b/modules/pytorch_wrapper/pto_adam.py @@ -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)) diff --git a/modules/pytorch_wrapper/pto_adamw.py b/modules/pytorch_wrapper/pto_adamw.py new file mode 100644 index 0000000..3f901fb --- /dev/null +++ b/modules/pytorch_wrapper/pto_adamw.py @@ -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,) \ No newline at end of file diff --git a/modules/pytorch_wrapper/pto_sgd.py b/modules/pytorch_wrapper/pto_sgd.py new file mode 100644 index 0000000..4fd5e4f --- /dev/null +++ b/modules/pytorch_wrapper/pto_sgd.py @@ -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,) \ No newline at end of file diff --git a/modules/pytorch_wrapper/pto_simple.py b/modules/pytorch_wrapper/pto_simple.py index ba7da8e..8ff50be 100644 --- a/modules/pytorch_wrapper/pto_simple.py +++ b/modules/pytorch_wrapper/pto_simple.py @@ -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 diff --git a/pyproject.toml b/pyproject.toml index ef88a29..3b30382 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.10" +version = "3.0.11" license = {file = "LICENSE"} [project.urls] diff --git a/tests/pytorch_wrapper/pto_adamw_test.py b/tests/pytorch_wrapper/pto_adamw_test.py new file mode 100644 index 0000000..0b2bb98 --- /dev/null +++ b/tests/pytorch_wrapper/pto_adamw_test.py @@ -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() diff --git a/tests/pytorch_wrapper/pto_sgd_test.py b/tests/pytorch_wrapper/pto_sgd_test.py new file mode 100644 index 0000000..973ddd9 --- /dev/null +++ b/tests/pytorch_wrapper/pto_sgd_test.py @@ -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()