Added lr scheduler nodes and cifar10 workflows.

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))

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,)

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,)

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,)