Added Pt Train Model node to allow specifying a Loss node for training.

__init__.py

@@ -109,6 +109,7 @@ from .modules.pytorch_wrapper.pt_to_rgb_tensors import PtToRgbTensors
 from .modules.pytorch_wrapper.pt_to_uint8 import PtToUint8
 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_train_model import PtTrainModel
 from .modules.pytorch_wrapper.pt_train_regression_model import PtTrainRegressionModel
 from .modules.pytorch_wrapper.pt_unsqueeze import PtUnsqueeze
 from .modules.pytorch_wrapper.pt_var import PtVar
@@ -300,6 +301,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
     "PtToUint8": PtToUint8,
     "PtTrainClassificationModel": PtTrainClassificationModel,
     "PtTrainClassificationModelLr": PtTrainClassificationModelLr,
+    "PtTrainModel": PtTrainModel,
     "PtTrainRegressionModel": PtTrainRegressionModel,
     "PtUnsqueeze": PtUnsqueeze,
     "PtVar": PtVar,
@@ -491,6 +493,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
     "PtToUint8": "Pt To Uint8",
     "PtTrainClassificationModel": "Pt Train Classification Model",
     "PtTrainClassificationModelLr": "Pt Train Classification Model Lr",
+    "PtTrainModel": "Pt Train Model",
     "PtTrainRegressionModel": "Pt Train Regression Model",
     "PtUnsqueeze": "Pt Unsqueeze",
     "PtVar": "Pt Var",

docs/model_training.md

@@ -95,3 +95,7 @@ Below is an example workflow for classifying the classic Iris dataset:
 `examples/workflows/iris_classification_train_eval.json`.
 
 According to [the UCI repository](https://archive.ics.uci.edu/dataset/53/iris), the neural network baseline accuracy for this dataset ranges from 92.105% to 100%. Using this workflow, you can expect around 96% accuracy, occasionally reaching 100%, which aligns with the baseline.
+
+# Specifying a loss function for training
+If you want to specify a loss function for training instead of using the default one in the trainer, use the `Pt Train Model` node. Create a loss function node (e.g., `Pt Cross Entropy Loss`) and connect it to the trainer node.
+Refer to `examples/workflows/train_model_with_custom_loss_example.json` for an example. (This workflow is intended to demonstrate node connections and is not optimized for high CIFAR10 accuracy. It runs for only 10 epochs to allow a quick check and reaches 88.9% accuracy.) You should adjust the hyperparameters according to your specific goals.

docs/reference/node_reference.md

@@ -225,6 +225,7 @@
 | [Pt Predict Regression Model](pt_predict_regression_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.   |
+| [Pt Train Model](pt_train_model.md) | Trains a model using a given dataset, loss function, optimizer, and number of epochs with learning rate decay.   |
 | [Pt Train Regression Model](pt_train_regression_model.md) | Trains a regression model using a given dataset, optimizer, and number of epochs.   |
 | [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. |

docs/reference/pt_train_model.md

@@ -0,0 +1,34 @@
+# Pt Train Model
+Trains a model using a given dataset, loss function, 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.
+
+Functionality-wise, this node is identical with PtTrainClassificationModelLr except that you can specify a custom loss function in this model where PtTrainClassificationModelLr uses nn.CrossEntropyLoss().
+
+## Input
+| Name | Data type |
+|---|---|
+| model | Ptmodel |
+| train_loader | Ptdataloader |
+| optimizer | Ptoptimizer |
+| loss_function | Ptloss |
+| epochs | Int |
+| use_gpu | Boolean |
+| early_stopping | Boolean |
+| early_stopping_rounds | Int |
+| output_best_val_model | Boolean |
+| scheduler | Ptlrscheduler |
+| val_loader | Ptdataloader |
+
+## Output
+| Data type |
+|---|
+| Ptmodel |
+| Tensor |
+| Tensor |
+
+<HR>
+Category: PyTorch wrapper - Training
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

modules/pytorch_wrapper/pt_train_classification_model_lr.py

@@ -100,7 +100,7 @@ class PtTrainClassificationModelLr:
             if use_gpu:
                 model.to("cuda")
 
-            loss_function = nn.CrossEntropyLoss()  # mean loss
+            loss_function = nn.CrossEntropyLoss()
 
             # Train the model
             for epoch in range(epochs):

modules/pytorch_wrapper/pt_train_model.py

@@ -0,0 +1,183 @@
+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
+
+
+class PtTrainModel:
+    """
+    Trains a model using a given dataset, loss function, 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.
+
+    Functionality-wise, this node is identical with PtTrainClassificationModelLr except that you can specify a custom loss function in this model where PtTrainClassificationModelLr uses nn.CrossEntropyLoss().
+
+    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", {}),
+                "loss_function": ("PTLOSS", {}),
+                "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": {
+                "scheduler": ("PTLRSCHEDULER", {}),
+                "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, 
+          loss_function: nn.Module,
+          epochs: int, 
+          use_gpu: bool, 
+          early_stopping: bool,
+          early_stopping_rounds: int,
+          output_best_val_model:bool,
+          scheduler: Optional[Union[torch.optim.lr_scheduler._LRScheduler, torch.optim.lr_scheduler.ReduceLROnPlateau]] = None,
+          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.
+            loss (torch.nn.Module): The loss class.
+            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).
+            scheduler (Optional[Union[torch.optim.lr_scheduler._LRScheduler, torch.optim.lr_scheduler.ReduceLROnPlateau]]): Learning rate scheduler.
+
+        Returns:
+            Tuple[torch.nn.Module, torch.Tensor, torch.Tensor]: A tuple containing the trained model, 
+            training loss history, and validation loss history.
+        """
+        if (scheduler is not None and 
+            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")
+
+            # 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 scheduler:
+                    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))

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.2.0"
+version = "3.2.1"
 license = {file = "LICENSE"}
 
 [project.urls]

tests/pytorch_wrapper/pt_train_model_test.py

@@ -0,0 +1,80 @@
+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_model import PtTrainModel
+from pytorch_wrapper.pt_save_model import PtSaveModel
+from pytorch_wrapper.pto_lr_scheduler_step import PtoLrSchedulerStep
+from pytorch_wrapper.ptn_cross_entropy_loss import PtnCrossEntropyLoss
+
+
+class TestPtTrainModel(unittest.TestCase):
+    def setUp(self):
+        """Set up test instance."""
+        self.loss_fun = PtnCrossEntropyLoss().f(reduction="mean")[0]
+        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 = PtTrainModel()
+        self.save_model = PtSaveModel()
+
+    def test_1(self):
+        epochs = 4
+        trained_model = self.trainer.f(
+                        self.model,
+                        self.train_loader,
+                        self.optimizer,
+                        self.loss_fun,
+                        epochs, # epochs
+                        use_gpu=True,
+                        early_stopping=False,
+                        early_stopping_rounds=1,
+                        scheduler=self.scheduler,
+                        output_best_val_model=False)[0]
+        self.save_model.f(trained_model, f"fashion_mnist_{epochs}_epochs_conv.pt")
+
+
+if __name__ == "__main__":
+    unittest.main()