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 from sklearn.metrics import accuracy_score from .rnn_output_reshape import extract_rnn_return_value_with_adjusted_label class PtTrainRNNModel: """ Pt Train RNN Model: Trains an RNN model using a given dataset, loss function, optimizer, and number of epochs with learning rate decay. Args: model (PTMODEL): The PyTorch model to be trained. linear_head (bool): The RNN model has a Linear head. train_loader (torch.utils.data.DataLoader): DataLoader containing the training dataset. optimizer (PTOPTIMIZER): Optimizer used for training. loss (PTLOSS): 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. classification_metrics (bool): If True and if val_loader is set, print out classification metrics. This is only valid for classification model training. use_valid_token_mean (bool): Computes the mean of the RNN output for non-zero inputs and feeds to the linear layer. val_loader (PTDATALOADER): DataLoader for validation dataset (optional). scheduler (PTLRSCHEDULER): Learning rate scheduler. h_0 (TENSOR): Optional initial hidden state. For unidirectional models, the shape is [num_layers, batch_size, hidden_size]. For bidirectional models, it is [2 * num_layers, batch_size, hidden_size]. If not provided, a zero tensor is used. 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", {}), "linear_head": ("BOOLEAN", {"default":True}), "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}), "classification_metrics": ("BOOLEAN", {"default":False}), "use_valid_token_mean" : ("BOOLEAN", {"default":True}) }, "optional": { "scheduler": ("PTLRSCHEDULER", {}), "val_loader": ("PTDATALOADER", {}), "h_0": ("TENSOR", {}), } } 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, linear_head: bool, 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, classification_metrics:bool, use_valid_token_mean: bool, scheduler: Optional[Union[torch.optim.lr_scheduler._LRScheduler, torch.optim.lr_scheduler.ReduceLROnPlateau]] = None, val_loader: Optional[torch.utils.data.DataLoader] = None, h_0: Optional[torch.Tensor] = None) -> Tuple[nn.Module, torch.Tensor, torch.Tensor]: """ Trains an RNN model. Args: model (torch.nn.Module): The PyTorch model to be trained. linear_head (bool): The RNN model has a Linear head. 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. classification_metrics (bool): If True and if val_loader is set, print out classification metrics. This is only valid for classification model training. use_valid_token_mean (bool): Computes the mean of the RNN output for non-zero inputs and feeds to the linear layer. 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. h_0 (optional[torch.Tensor]): Optional initial hidden state. For unidirectional models, the shape is [num_layers, batch_size, hidden_size]. For bidirectional models, it is [2 * num_layers, batch_size, hidden_size]. If not provided, a zero tensor is used. 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 x, y in train_loader: if use_gpu: if isinstance(x, list): x = [e.to("cuda") for e in x] else: x = x.to("cuda") y = y.to("cuda") optimizer.zero_grad() if linear_head: if h_0: y_hat = model(x, hx=h_0) # Note that both RNN and GRU use hx. See PyTorch code. https://github.com/pytorch/pytorch/blob/v2.6.0/torch/nn/modules/rnn.py else: y_hat = model(x) if y.dim() == 3: y = y[:,-1,:] # Grab last token elif y.dim() == 1: y = torch.unsqueeze(y, dim=-1) y = y.to(torch.float32) else: if h_0: output_seq, _ = model(x, hx=h_0) else: output_seq, _ = model(x) outputs, adjusted_labels = extract_rnn_return_value_with_adjusted_label( x, y, output_seq, bidirectional=model.bidirectional, return_valid_token_mean=use_valid_token_mean ) if outputs.size() != adjusted_labels.size(): raise ValueError(f"y_hat size : {outputs.size()} does not match label size: {adjusted_labels.size()}") y = adjusted_labels y_hat = outputs loss = loss_function(y_hat, y) loss.backward() optimizer.step() batch_size = y.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(): all_preds = [] all_labels = [] 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 x, y in val_loader: if use_gpu: if isinstance(x, list): x = [e.to("cuda") for e in x] else: x = x.to("cuda") y = y.to("cuda") if linear_head: y_hat = model(x) if y.dim() == 3: y = y[:,-1,:] # Grab last token elif y.dim() == 1: y = torch.unsqueeze(y, dim=-1) y = y.to(torch.float32) else: output_seq, _ = model(x) outputs, adjusted_labels = extract_rnn_return_value_with_adjusted_label( x, y, output_seq, bidirectional=model.bidirectional, return_valid_token_mean=use_valid_token_mean ) if outputs.size() != adjusted_labels.size(): raise ValueError(f"y_hat size : {outputs.size()} does not match label size: {adjusted_labels.size()}") y = adjusted_labels y_hat = outputs loss = loss_function(y_hat, y) batch_size = y.size(0) batch_loss = loss.item() * batch_size total_loss += batch_loss if classification_metrics and y.squeeze().dim() == 1 : # Accuracy collection for binary classification preds = (F.sigmoid(y_hat) > 0.5).int().detach().cpu().numpy() labels = y.int().detach().cpu().numpy() all_preds.extend(preds) all_labels.extend(labels) epoch_loss = total_loss / total_samples print(f"Epoch (val) {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}") val_loss.append(epoch_loss) if classification_metrics: if len(all_labels): # Compute accuracy acc = accuracy_score(all_labels, all_preds) print(f"Validation Accuracy: {acc:.4f}") else: print("Accuracy was not computed as the label is not scalar.") 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))