diff --git a/__init__.py b/__init__.py index 8e833cc..e40fc7b 100644 --- a/__init__.py +++ b/__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", diff --git a/docs/model_training.md b/docs/model_training.md index adbc5b5..84e52d4 100644 --- a/docs/model_training.md +++ b/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. \ No newline at end of file diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index 0be6173..a050511 100644 --- a/docs/reference/node_reference.md +++ b/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. | diff --git a/docs/reference/pt_train_model.md b/docs/reference/pt_train_model.md new file mode 100644 index 0000000..104ddfc --- /dev/null +++ b/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 | + +
+Category: PyTorch wrapper - Training + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/examples/workflows/train_model_with_custom_loss_example.json b/examples/workflows/train_model_with_custom_loss_example.json new file mode 100644 index 0000000..b49d2dd --- /dev/null +++ b/examples/workflows/train_model_with_custom_loss_example.json @@ -0,0 +1 @@ +{"last_node_id":95,"last_link_id":119,"nodes":[{"id":61,"type":"PtvDataset","pos":[2007.994873046875,1723.90625],"size":[400,200],"flags":{},"order":10,"mode":0,"inputs":[{"name":"transform","localized_name":"transform","type":"PTVTRANSFORM","shape":7,"link":103}],"outputs":[{"name":"PTVDATASET","localized_name":"PTVDATASET","type":"PTVDATASET","links":[74,83],"slot_index":0}],"properties":{"Node name for S&R":"PtvDataset"},"widgets_values":["CIFAR10",true,"","{\"train\":True}"]},{"id":19,"type":"PtSaveModel","pos":[3530.668701171875,1647.40185546875],"size":[396.8008117675781,78],"flags":{},"order":21,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":119},{"name":"model_path","type":"STRING","widget":{"name":"model_path"},"link":99}],"outputs":[{"name":"model full path","localized_name":"model full path","type":"STRING","links":[70],"slot_index":0}],"properties":{"Node name for S&R":"PtSaveModel"},"widgets_values":["/container/transfer/models/cifar10.pt"]},{"id":62,"type":"PtvDataset","pos":[2007.994873046875,2039.8226318359375],"size":[400,200],"flags":{},"order":11,"mode":0,"inputs":[{"name":"transform","localized_name":"transform","type":"PTVTRANSFORM","shape":7,"link":89}],"outputs":[{"name":"PTVDATASET","localized_name":"PTVDATASET","type":"PTVDATASET","links":[75,86],"slot_index":0}],"properties":{"Node name for S&R":"PtvDataset"},"widgets_values":["CIFAR10",true,"","{\"train\":False}"]},{"id":76,"type":"PtvDatasetLen","pos":[2007.994873046875,1638.4072265625],"size":[210,26],"flags":{},"order":14,"mode":0,"inputs":[{"name":"dataset","localized_name":"dataset","type":"PTVDATASET","link":83}],"outputs":[{"name":"INT","localized_name":"INT","type":"INT","links":[84],"slot_index":0}],"properties":{"Node name for S&R":"PtvDatasetLen"},"widgets_values":[]},{"id":78,"type":"PtvDatasetLen","pos":[2007.994873046875,2297.724365234375],"size":[210,26],"flags":{},"order":16,"mode":0,"inputs":[{"name":"dataset","localized_name":"dataset","type":"PTVDATASET","link":86}],"outputs":[{"name":"INT","localized_name":"INT","type":"INT","links":[85],"slot_index":0}],"properties":{"Node name for S&R":"PtvDatasetLen"},"widgets_values":[]},{"id":79,"type":"CDAShowInt","pos":[2007.994873046875,2381.43359375],"size":[315,112],"flags":{},"order":19,"mode":0,"inputs":[{"name":"int_scalar","type":"INT","widget":{"name":"int_scalar"},"link":85}],"outputs":[{"name":"STRING","localized_name":"STRING","type":"STRING","shape":6,"links":null}],"properties":{"Node name for S&R":"CDAShowInt"},"widgets_values":[0,"10000"]},{"id":53,"type":"MarkdownNote","pos":[1215.7994384765625,1723.90625],"size":[210,88],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Load train data"],"color":"#432","bgcolor":"#653"},{"id":54,"type":"MarkdownNote","pos":[1214.4312744140625,2039.8226318359375],"size":[210,88],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Load val data"],"color":"#432","bgcolor":"#653"},{"id":77,"type":"CDAShowInt","pos":[2007.994873046875,1504.9381103515625],"size":[315,112],"flags":{},"order":18,"mode":0,"inputs":[{"name":"int_scalar","type":"INT","widget":{"name":"int_scalar"},"link":84}],"outputs":[{"name":"STRING","localized_name":"STRING","type":"STRING","shape":6,"links":null}],"properties":{"Node name for S&R":"CDAShowInt"},"widgets_values":[0,"50000"]},{"id":90,"type":"PtvTransformsDataAugment","pos":[1595.48486328125,1723.90625],"size":[340.20001220703125,202],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[{"name":"PTVTRANSFORM","localized_name":"PTVTRANSFORM","type":"PTVTRANSFORM","links":[103],"slot_index":0}],"properties":{"Node name for S&R":"PtvTransformsDataAugment"},"widgets_values":[0.5,0,0,0.125,0.125,1,1]},{"id":81,"type":"PtvTransformsToTensor","pos":[1595.48486328125,2039.8226318359375],"size":[338.8136291503906,26],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[{"name":"PTVTRANSFORM","localized_name":"PTVTRANSFORM","type":"PTVTRANSFORM","links":[89],"slot_index":0}],"properties":{"Node name for S&R":"PtvTransformsToTensor"},"widgets_values":[]},{"id":56,"type":"MarkdownNote","pos":[3535.058349609375,1775.04931640625],"size":[210,88],"flags":{},"order":4,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Save model"],"color":"#432","bgcolor":"#653"},{"id":80,"type":"PtnResnetModel","pos":[2478.35009765625,1413.751953125],"size":[379.3387145996094,113.2640609741211],"flags":{},"order":5,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[91,113],"slot_index":0}],"properties":{"Node name for S&R":"PtnResnetModel"},"widgets_values":["(3,32,32)",10,10]},{"id":63,"type":"PtDataLoader","pos":[2473.048828125,1723.90625],"size":[400,200],"flags":{},"order":13,"mode":0,"inputs":[{"name":"dataset","localized_name":"dataset","type":"PTVDATASET","link":74}],"outputs":[{"name":"PTDATALOADER","localized_name":"PTDATALOADER","type":"PTDATALOADER","links":[114],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[128,true,"{\"num_workers\":1}"]},{"id":11,"type":"PtoAdam","pos":[2976.371337890625,1387.1434326171875],"size":[246.16842651367188,110.40398406982422],"flags":{},"order":12,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":91}],"outputs":[{"name":"PTOPTIMIZER","localized_name":"PTOPTIMIZER","type":"PTOPTIMIZER","links":[111,115],"slot_index":0}],"properties":{"Node name for S&R":"PtoAdam"},"widgets_values":[0.001,0.9,0.999]},{"id":64,"type":"PtDataLoader","pos":[2473.048828125,2039.8226318359375],"size":[400,200],"flags":{},"order":15,"mode":0,"inputs":[{"name":"dataset","localized_name":"dataset","type":"PTVDATASET","link":75}],"outputs":[{"name":"PTDATALOADER","localized_name":"PTDATALOADER","type":"PTDATALOADER","links":[117],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[128,false,"{\"num_workers\":1}"]},{"id":95,"type":"PtnCrossEntropyLoss","pos":[2481.095947265625,1593.5537109375],"size":[368.3888244628906,61.8134880065918],"flags":{},"order":6,"mode":0,"inputs":[],"outputs":[{"name":"PTLOSS","localized_name":"PTLOSS","type":"PTLOSS","links":[116],"slot_index":0}],"properties":{"Node name for S&R":"PtnCrossEntropyLoss"},"widgets_values":["mean"]},{"id":93,"type":"PtoLrSchedulerCosineAnnealing","pos":[3274.407958984375,1388.4407958984375],"size":[280.896240234375,86.71690368652344],"flags":{},"order":17,"mode":0,"inputs":[{"name":"optimizer","localized_name":"optimizer","type":"PTOPTIMIZER","link":111}],"outputs":[{"name":"PTLRSCHEDULER","localized_name":"PTLRSCHEDULER","type":"PTLRSCHEDULER","links":[118],"slot_index":0}],"properties":{"Node name for S&R":"PtoLrSchedulerCosineAnnealing"},"widgets_values":[10,0]},{"id":51,"type":"PandasShowText","pos":[3540.85546875,1921.258056640625],"size":[333.6095275878906,112],"flags":{},"order":22,"mode":0,"inputs":[{"name":"text","type":"STRING","widget":{"name":"text"},"link":70}],"outputs":[{"name":"STRING","localized_name":"STRING","type":"STRING","shape":6,"links":null}],"properties":{"Node name for S&R":"PandasShowText"},"widgets_values":["","/container/transfer/models/cifar10_tmp.pt"]},{"id":55,"type":"MarkdownNote","pos":[2988.48876953125,1922.287353515625],"size":[232.82992553710938,88],"flags":{},"order":7,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Train"],"color":"#432","bgcolor":"#653"},{"id":52,"type":"MarkdownNote","pos":[2484.47021484375,1271.1656494140625],"size":[210,88],"flags":{},"order":8,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["# Define model"],"color":"#432","bgcolor":"#653"},{"id":89,"type":"CDAStringCreate","pos":[3536.05712890625,1534.6680908203125],"size":[385.5651550292969,59.037696838378906],"flags":{},"order":9,"mode":0,"inputs":[],"outputs":[{"name":"STRING","localized_name":"STRING","type":"STRING","links":[99],"slot_index":0}],"properties":{"Node name for S&R":"CDAStringCreate"},"widgets_values":["/container/transfer/models/cifar10_tmp.pt"]},{"id":94,"type":"PtTrainModel","pos":[2985.846923828125,1603.26904296875],"size":[315,254],"flags":{},"order":20,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":113},{"name":"train_loader","localized_name":"train_loader","type":"PTDATALOADER","link":114},{"name":"optimizer","localized_name":"optimizer","type":"PTOPTIMIZER","link":115},{"name":"loss_function","localized_name":"loss_function","type":"PTLOSS","link":116},{"name":"scheduler","localized_name":"scheduler","type":"PTLRSCHEDULER","shape":7,"link":118},{"name":"val_loader","localized_name":"val_loader","type":"PTDATALOADER","shape":7,"link":117}],"outputs":[{"name":"Model","localized_name":"Model","type":"PTMODEL","links":[119],"slot_index":0},{"name":"train loss","localized_name":"train loss","type":"TENSOR","links":null},{"name":"val loss","localized_name":"val loss","type":"TENSOR","links":null}],"properties":{"Node name for S&R":"PtTrainModel"},"widgets_values":[10,true,false,10,false]}],"links":[[70,19,0,51,0,"STRING"],[74,61,0,63,0,"PTVDATASET"],[75,62,0,64,0,"PTVDATASET"],[83,61,0,76,0,"PTVDATASET"],[84,76,0,77,0,"INT"],[85,78,0,79,0,"INT"],[86,62,0,78,0,"PTVDATASET"],[89,81,0,62,0,"PTVTRANSFORM"],[91,80,0,11,0,"PTMODEL"],[99,89,0,19,1,"STRING"],[103,90,0,61,0,"PTVTRANSFORM"],[111,11,0,93,0,"PTOPTIMIZER"],[113,80,0,94,0,"PTMODEL"],[114,63,0,94,1,"PTDATALOADER"],[115,11,0,94,2,"PTOPTIMIZER"],[116,95,0,94,3,"PTLOSS"],[117,64,0,94,5,"PTDATALOADER"],[118,93,0,94,4,"PTLRSCHEDULER"],[119,94,0,19,0,"PTMODEL"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.9598170997212797,"offset":[-2396.5953376811794,-1204.7438265178134]},"node_versions":{"pt-wrapper":"f42fef60b92ca64335192e9c97b23fe85f23b0ad","data-analysis":"68dd090367adffa569c3b3993d56e422f4516d2d"},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/modules/pytorch_wrapper/pt_train_classification_model_lr.py b/modules/pytorch_wrapper/pt_train_classification_model_lr.py index 2dcf095..38081a2 100644 --- a/modules/pytorch_wrapper/pt_train_classification_model_lr.py +++ b/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): diff --git a/modules/pytorch_wrapper/pt_train_model.py b/modules/pytorch_wrapper/pt_train_model.py new file mode 100644 index 0000000..c8735b0 --- /dev/null +++ b/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)) \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index eb83abe..365f883 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.2.0" +version = "3.2.1" license = {file = "LICENSE"} [project.urls] diff --git a/tests/pytorch_wrapper/pt_train_model_test.py b/tests/pytorch_wrapper/pt_train_model_test.py new file mode 100644 index 0000000..1ed1cd2 --- /dev/null +++ b/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()