Added support for RNN
This commit is contained in:
HowToSD
@@ -0,0 +1,129 @@
|
||||
from typing import Tuple
|
||||
import re
|
||||
from html import unescape
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
from datasets import load_dataset # Hugging Face (not PyTorch)
|
||||
import scipy
|
||||
import numpy as np
|
||||
# Monkey-patch gensim as it relies on scipy's deprecated triu method
|
||||
# /gensim/matutils.py has this line which causes ImportError:
|
||||
# from scipy.linalg import get_blas_funcs, triu
|
||||
# TODO: Monitor gensim and fix this if upgrade is made
|
||||
scipy.linalg.triu = np.triu
|
||||
import gensim.downloader as api
|
||||
|
||||
class HfGloveDataset(Dataset):
|
||||
"""
|
||||
Hugging Face Dataset wrapper with word embedding support.
|
||||
|
||||
pragma: skip_doc
|
||||
"""
|
||||
def __init__(self,
|
||||
dataset_name: str,
|
||||
split: str,
|
||||
glove_dim: int = 100,
|
||||
max_seq_len: int = 256,
|
||||
sample_field_name: str = "text",
|
||||
label_field_name: str = "label") -> None:
|
||||
"""
|
||||
Initializes the dataset by loading Hugging Face dataset and GloVe vectors.
|
||||
|
||||
Args:
|
||||
dataset_name (str): The name of the Hugging Face dataset.
|
||||
split (str): The dataset split (e.g., 'train').
|
||||
glove_dim (int): Dimensionality of GloVe embeddings.
|
||||
max_seq_len (int): Maximum sequence length.
|
||||
sample_field_name (str): Field name for text samples.
|
||||
label_field_name (str): Field name for labels.
|
||||
"""
|
||||
self.dataset = load_dataset(dataset_name, split=split)
|
||||
self.glove_dim = glove_dim
|
||||
self.max_seq_len = max_seq_len
|
||||
self.sample_field_name = sample_field_name
|
||||
self.label_field_name = label_field_name
|
||||
|
||||
# Load GloVe from gensim
|
||||
glove_name = f"glove-wiki-gigaword-{glove_dim}"
|
||||
self.glove = api.load(glove_name)
|
||||
self.unk_vector = torch.zeros(glove_dim)
|
||||
|
||||
def tokenize(self, text: str) -> list[str]:
|
||||
"""
|
||||
Tokenizes raw text into a list of lowercase word tokens.
|
||||
It removes HTML tags, decodes HTML entities, strips punctuation
|
||||
(while keeping letters and numbers), and splits on whitespace.
|
||||
|
||||
Args:
|
||||
text (str): Raw text input (may include HTML and punctuation).
|
||||
|
||||
Returns:
|
||||
list[str]: List of clean, lowercase word tokens.
|
||||
"""
|
||||
# Decode HTML entities like "&", "<", etc. to their actual characters
|
||||
text = unescape(text)
|
||||
|
||||
# Remove HTML tags (like <br />, <p>, etc.)
|
||||
# The regex "<.*?>" uses .*? (non-greedy match) to ensure it matches one tag at a time,
|
||||
# instead of greedily removing everything between the first '<' and last '>'
|
||||
text = re.sub(r"<.*?>", " ", text)
|
||||
|
||||
# Remove all characters except letters, digits, and whitespace
|
||||
# This strips punctuation like "!", ".", "(", but keeps words and numbers
|
||||
text = re.sub(r"[^a-zA-Z0-9\s]", " ", text)
|
||||
|
||||
# Convert to lowercase and split by whitespace to get token list
|
||||
return text.lower().split()
|
||||
|
||||
|
||||
def vectorize(self, tokens: list[str]) -> torch.Tensor:
|
||||
"""
|
||||
Converts tokens to a fixed-length tensor of word vectors.
|
||||
|
||||
Args:
|
||||
tokens (list[str]): A list of word tokens.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: A tensor of shape (max_seq_len, glove_dim).
|
||||
"""
|
||||
vectors = []
|
||||
for token in tokens[:self.max_seq_len]:
|
||||
if token in self.glove:
|
||||
vec = torch.tensor(self.glove[token], dtype=torch.float32)
|
||||
else:
|
||||
vec = self.unk_vector
|
||||
vectors.append(vec)
|
||||
|
||||
if len(vectors) < self.max_seq_len:
|
||||
pad = [torch.zeros(self.glove_dim) for _ in range(self.max_seq_len - len(vectors))]
|
||||
vectors.extend(pad)
|
||||
|
||||
return torch.stack(vectors)
|
||||
|
||||
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Retrieves the (embedded sequence, label) pair at the given index.
|
||||
|
||||
Args:
|
||||
idx (int): Index of the item.
|
||||
|
||||
Returns:
|
||||
Tuple[torch.Tensor, torch.Tensor]: Embedding tensor and label tensor.
|
||||
"""
|
||||
# Get a text sample from database and convert to a word embedding
|
||||
sample = self.dataset[idx]
|
||||
tokens = self.tokenize(sample[self.sample_field_name])
|
||||
emb_tensor = self.vectorize(tokens)
|
||||
|
||||
# Get a corresponding label
|
||||
label = torch.tensor(sample[self.label_field_name], dtype=torch.long)
|
||||
return emb_tensor, label
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""
|
||||
Returns the number of samples in the dataset.
|
||||
|
||||
Returns:
|
||||
int: Dataset length.
|
||||
"""
|
||||
return len(self.dataset)
|
||||
@@ -65,7 +65,7 @@ class PtTrainModel:
|
||||
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.
|
||||
Trains a model.
|
||||
|
||||
Args:
|
||||
model (torch.nn.Module): The PyTorch model to be trained.
|
||||
@@ -117,6 +117,7 @@ class PtTrainModel:
|
||||
labels = labels.to("cuda")
|
||||
optimizer.zero_grad()
|
||||
outputs = model(images)
|
||||
|
||||
loss = loss_function(outputs, labels)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
@@ -0,0 +1,262 @@
|
||||
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:
|
||||
"""
|
||||
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.
|
||||
|
||||
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", {}),
|
||||
}
|
||||
}
|
||||
|
||||
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) -> 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.
|
||||
|
||||
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
|
||||
|
||||
all_preds = []
|
||||
all_labels = []
|
||||
|
||||
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:
|
||||
x = x.to("cuda")
|
||||
y = y.to("cuda")
|
||||
optimizer.zero_grad()
|
||||
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)
|
||||
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():
|
||||
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:
|
||||
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 = (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))
|
||||
@@ -0,0 +1,116 @@
|
||||
from typing import Any, Dict, Tuple
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class PtnRNN:
|
||||
"""
|
||||
Ptn RNN:
|
||||
A recurrent neural network (RNN) model consisting of one or more of a recurrent layer.
|
||||
|
||||
Args:
|
||||
input_size (int): The number of input features.
|
||||
hidden_size (int): The number of output features of the hidden layer matrix.
|
||||
num_layers (int): Number of hidden layers.
|
||||
nonlinearity (str): Activation function to apply after each hidden layer. Specify 'tanh' or 'relu'.
|
||||
bias (bool): Use bias or not.
|
||||
batch_first (bool): Indicates the shape of the input. Input is assumed to be a rank-3 tensor. If True, the shape is [Batch, Seq, Token]. If False, the shape is [Seq, Batch, Token]. Note that default is set to `True` unlike PyTorch's RNN default model parameter.
|
||||
dropout (float): Dropout probability to be applied to the output of intermediate layers if non-zero. This will not be applied to the last layer's output.
|
||||
bidirectional (bool): Processes the input in a backward direction and append the result to the output of the forward direction.
|
||||
|
||||
category: PyTorch wrapper - Model
|
||||
"""
|
||||
|
||||
@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": {
|
||||
"input_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
|
||||
"hidden_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
|
||||
"num_layers": ("INT", {"default": 1, "min": 1, "max": 1e3}),
|
||||
"nonlinearity": (("tanh", "relu"),),
|
||||
"bias": ("BOOLEAN", {"default": True}),
|
||||
"batch_first": ("BOOLEAN", {"default": True}),
|
||||
"dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1}),
|
||||
"bidirectional": ("BOOLEAN", {"default": False}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES: Tuple[str] = ("PTMODEL",)
|
||||
FUNCTION: str = "f"
|
||||
CATEGORY: str = "Training"
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kw):
|
||||
return float("NaN")
|
||||
|
||||
def f(self,
|
||||
input_size: int,
|
||||
hidden_size: int,
|
||||
num_layers: int,
|
||||
nonlinearity: str,
|
||||
bias: bool,
|
||||
batch_first: bool,
|
||||
dropout: float,
|
||||
bidirectional: bool
|
||||
) -> Tuple[nn.Module]:
|
||||
"""
|
||||
Constructs an RNN model.
|
||||
|
||||
Args:
|
||||
input_size (int): The number of input features.
|
||||
hidden_size (int): The number of output features of the hidden layer matrix.
|
||||
num_layers (int): Number of hidden layers.
|
||||
nonlinearity (str): Activation function to apply after each hidden layer. Specify 'tanh' or 'relu'.
|
||||
bias (bool): Use bias or not.
|
||||
batch_first (bool): Indicates the shape of the input. Input is assumed to be a rank-3 tensor. If True, the shape is [Batch, Seq, Token]. If False, the shape is [Seq, Batch, Token].
|
||||
dropout (float): Dropout probability to be applied to the output of intermediate layers if non-zero. This will not be applied to the last layer's output.
|
||||
bidirectional (bool): Processes the input in a backward direction and append the result to the output of the forward direction.
|
||||
Returns:
|
||||
Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
|
||||
"""
|
||||
with torch.inference_mode(False):
|
||||
model = nn.RNN(
|
||||
input_size=input_size,
|
||||
hidden_size=hidden_size,
|
||||
num_layers=num_layers,
|
||||
nonlinearity=nonlinearity,
|
||||
bias=bias,
|
||||
batch_first=batch_first,
|
||||
dropout=dropout,
|
||||
bidirectional=bidirectional
|
||||
)
|
||||
|
||||
directions = 2 if bidirectional else 1
|
||||
|
||||
# Initialize weights
|
||||
for layer in range(num_layers):
|
||||
for direction in range(directions):
|
||||
suffix = f"_reverse" if direction == 1 else ""
|
||||
|
||||
weight_ih = getattr(model, f"weight_ih_l{layer}{suffix}")
|
||||
weight_hh = getattr(model, f"weight_hh_l{layer}{suffix}")
|
||||
|
||||
# ih for X, hh is for prev H to current H
|
||||
if nonlinearity == "relu":
|
||||
torch.nn.init.kaiming_normal_(weight_ih, nonlinearity='relu')
|
||||
torch.nn.init.kaiming_normal_(weight_hh, nonlinearity='relu')
|
||||
else:
|
||||
torch.nn.init.xavier_normal_(weight_ih)
|
||||
torch.nn.init.xavier_normal_(weight_hh)
|
||||
|
||||
# bias_ih = getattr(model, f"bias_ih_l{layer}{suffix}")
|
||||
# bias_hh = getattr(model, f"bias_hh_l{layer}{suffix}")
|
||||
|
||||
# torch.nn.init.uniform_(weight_ih, -0.05, 0.05)
|
||||
# torch.nn.init.uniform_(weight_hh, -0.05, 0.05)
|
||||
# nn.init.zeros_(bias_ih)
|
||||
# nn.init.zeros_(bias_hh)
|
||||
|
||||
return (model,)
|
||||
@@ -0,0 +1,184 @@
|
||||
from typing import Any, Dict, Tuple
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from pytorch_wrapper.rnn_output_reshape import extract_rnn_return_value_with_adjusted_label
|
||||
|
||||
|
||||
class RNNLinear(nn.Module):
|
||||
"""
|
||||
A recurrent neural network (RNN) with a linear output layer.
|
||||
|
||||
This module uses a standard PyTorch RNN, followed by a linear layer applied to the final output tensor.
|
||||
|
||||
pragma: skip_doc
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
input_size: int,
|
||||
hidden_size: int,
|
||||
num_layers: int,
|
||||
nonlinearity: str,
|
||||
bias: bool,
|
||||
batch_first: bool,
|
||||
dropout: float,
|
||||
bidirectional: bool,
|
||||
linear_output_size: int,
|
||||
linear_bias: bool
|
||||
):
|
||||
super().__init__()
|
||||
self.rnn_model = nn.RNN(
|
||||
input_size=input_size,
|
||||
hidden_size=hidden_size,
|
||||
num_layers=num_layers,
|
||||
nonlinearity=nonlinearity,
|
||||
bias=bias,
|
||||
batch_first=batch_first,
|
||||
dropout=dropout,
|
||||
bidirectional=bidirectional
|
||||
)
|
||||
self.bidirectional = bidirectional
|
||||
self.batch_first = batch_first
|
||||
directions = 2 if bidirectional else 1
|
||||
|
||||
# Initialize RNN weights
|
||||
for layer in range(num_layers):
|
||||
for direction in range(directions):
|
||||
suffix = f"_reverse" if direction == 1 else ""
|
||||
weight_ih = getattr(self.rnn_model, f"weight_ih_l{layer}{suffix}")
|
||||
weight_hh = getattr(self.rnn_model, f"weight_hh_l{layer}{suffix}")
|
||||
|
||||
if nonlinearity == "relu":
|
||||
torch.nn.init.kaiming_normal_(weight_ih, nonlinearity='relu')
|
||||
torch.nn.init.kaiming_normal_(weight_hh, nonlinearity='relu')
|
||||
else:
|
||||
torch.nn.init.xavier_normal_(weight_ih)
|
||||
torch.nn.init.xavier_normal_(weight_hh)
|
||||
|
||||
self.linear_model = nn.Linear(
|
||||
in_features=hidden_size * directions,
|
||||
out_features=linear_output_size,
|
||||
bias=linear_bias
|
||||
)
|
||||
torch.nn.init.xavier_uniform_(self.linear_model.weight)
|
||||
|
||||
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Forward pass through the RNN and the linear layer.
|
||||
|
||||
Args:
|
||||
inputs (torch.Tensor): Input tensor of shape (batch, seq, input_size) if batch_first=True.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Output of the linear layer.
|
||||
"""
|
||||
out_seq, _ = self.rnn_model(inputs)
|
||||
|
||||
x = extract_rnn_return_value_with_adjusted_label(
|
||||
inputs, # This is needed to find out padded tokens
|
||||
None,
|
||||
out_seq,
|
||||
self.bidirectional,
|
||||
self.batch_first,
|
||||
False,
|
||||
return_valid_token_mean=True
|
||||
)
|
||||
|
||||
outputs = self.linear_model(x)
|
||||
return outputs
|
||||
|
||||
|
||||
class PtnRNNLinear:
|
||||
"""
|
||||
Ptn RNN Linear:
|
||||
A recurrent neural network (RNN) model with a linear head.
|
||||
|
||||
Args:
|
||||
input_size (int): The number of input features.
|
||||
hidden_size (int): The number of output features of the hidden layer matrix.
|
||||
num_layers (int): Number of hidden layers in RNN.
|
||||
nonlinearity (str): Activation function to apply after each hidden layer. Specify 'tanh' or 'relu'.
|
||||
bias (bool): Use bias or not.
|
||||
batch_first (bool): Indicates the shape of the input. Input is assumed to be a rank-3 tensor. If True, the shape is [Batch, Seq, Token]. If False, the shape is [Seq, Batch, Token]. Note that default is set to `True` unlike PyTorch's RNN default model parameter.
|
||||
dropout (float): Dropout probability to be applied to the output of intermediate layers if non-zero. This will not be applied to the last layer's output.
|
||||
bidirectional (bool): Processes the input in a backward direction and append the result to the output of the forward direction.
|
||||
linear_output_size (int): The number of output features of the linear layer.
|
||||
linar_bias (bool): Use bias or not in linear layer.
|
||||
|
||||
category: PyTorch wrapper - Model
|
||||
"""
|
||||
|
||||
@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": {
|
||||
"input_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
|
||||
"hidden_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
|
||||
"num_layers": ("INT", {"default": 1, "min": 1, "max": 1e3}),
|
||||
"nonlinearity": (("tanh", "relu"),),
|
||||
"bias": ("BOOLEAN", {"default": True}),
|
||||
"batch_first": ("BOOLEAN", {"default": True}),
|
||||
"dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1}),
|
||||
"bidirectional": ("BOOLEAN", {"default": False}),
|
||||
"linear_output_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
|
||||
"linear_bias": ("BOOLEAN", {"default": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES: Tuple[str] = ("PTMODEL",)
|
||||
FUNCTION: str = "f"
|
||||
CATEGORY: str = "Training"
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kw):
|
||||
return float("NaN")
|
||||
|
||||
def f(self,
|
||||
input_size: int,
|
||||
hidden_size: int,
|
||||
num_layers: int,
|
||||
nonlinearity: str,
|
||||
bias: bool,
|
||||
batch_first: bool,
|
||||
dropout: float,
|
||||
bidirectional: bool,
|
||||
linear_output_size: int,
|
||||
linear_bias: bool
|
||||
) -> Tuple[nn.Module]:
|
||||
"""
|
||||
Constructs an RNN model.
|
||||
|
||||
Args:
|
||||
input_size (int): The number of input features.
|
||||
hidden_size (int): The number of output features of the hidden layer matrix.
|
||||
num_layers (int): Number of hidden layers.
|
||||
nonlinearity (str): Activation function to apply after each hidden layer. Specify 'tanh' or 'relu'.
|
||||
bias (bool): Use bias or not.
|
||||
batch_first (bool): Indicates the shape of the input. Input is assumed to be a rank-3 tensor. If True, the shape is [Batch, Seq, Token]. If False, the shape is [Seq, Batch, Token].
|
||||
dropout (float): Dropout probability to be applied to the output of intermediate layers if non-zero. This will not be applied to the last layer's output.
|
||||
bidirectional (bool): Processes the input in a backward direction and append the result to the output of the forward direction.
|
||||
linear_output_size (int): The number of output features of the linear layer.
|
||||
linar_bias (bool): Use bias or not in linear layer.
|
||||
|
||||
Returns:
|
||||
Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
|
||||
"""
|
||||
with torch.inference_mode(False):
|
||||
model = RNNLinear(
|
||||
input_size=input_size,
|
||||
hidden_size=hidden_size,
|
||||
num_layers=num_layers,
|
||||
nonlinearity=nonlinearity,
|
||||
bias=bias,
|
||||
batch_first=batch_first,
|
||||
dropout=dropout,
|
||||
bidirectional=bidirectional,
|
||||
linear_output_size=linear_output_size,
|
||||
linear_bias=linear_bias
|
||||
)
|
||||
return (model,)
|
||||
@@ -0,0 +1,80 @@
|
||||
from typing import Any, Dict, Tuple
|
||||
import torch
|
||||
|
||||
|
||||
class PtvHfGloveDataset:
|
||||
"""
|
||||
A PyTorch Dataset class wrapper for a Hugging Face dataset that converts text to embedding using Glove.
|
||||
|
||||
Args:
|
||||
name (str): Name of the dataset class (e.g., MNIST, FashionMNIST).
|
||||
split (str): The name of the subset of dataset to retrieve or the split specifier.
|
||||
embed_dim (int): Embedding dimension for Glove.
|
||||
max_seq_len (int): Sequence length.
|
||||
sample_field_name (str): Field name for text samples.
|
||||
label_field_name (str): Field name for labels.
|
||||
|
||||
category: PyTorch wrapper - Training
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
||||
"""
|
||||
Defines the required input types.
|
||||
|
||||
Returns:
|
||||
Dict[str, Any]: A dictionary of required input types.
|
||||
"""
|
||||
return {
|
||||
"required": {
|
||||
"name": ("STRING", {"default": "", "multiline": False}),
|
||||
"split": ("STRING", {"default": "train","multiline": False}),
|
||||
"embed_dim": ("INT", {"default": 100, "min":32, "max": 4096}),
|
||||
"max_seq_len": ("INT", {"default": 256, "min":8, "max": 4096}),
|
||||
"sample_field_name": ("STRING", {"default": "text","multiline": False}),
|
||||
"label_field_name": ("STRING", {"default": "label","multiline": False}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES: Tuple[str] = ("PTVDATASET",)
|
||||
FUNCTION: str = "f"
|
||||
CATEGORY: str = "Training"
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kw):
|
||||
return float("NaN")
|
||||
|
||||
def f(self,
|
||||
name: str,
|
||||
split: str,
|
||||
embed_dim: int,
|
||||
max_seq_len: int,
|
||||
sample_field_name: str,
|
||||
label_field_name: str) -> Tuple:
|
||||
"""
|
||||
Loads a dataset from Hugging Face with specified parameters.
|
||||
|
||||
Args:
|
||||
name (str): Name of the dataset class (e.g., MNIST, FashionMNIST).
|
||||
split (str): The name of the subset of dataset to retrieve or the split specifier.
|
||||
embed_dim (int): Embedding dimension for Glove.
|
||||
max_seq_len (int): Sequence length.
|
||||
sample_field_name (str): Field name for text samples.
|
||||
label_field_name (str): Field name for labels.
|
||||
|
||||
Returns:
|
||||
Tuple: A tuple containing the dataset instance.
|
||||
"""
|
||||
# Keep import within this method so that relevant packages are imported
|
||||
# only when they are actually needed.
|
||||
from .hf_glove_dataset import HfGloveDataset
|
||||
|
||||
with torch.inference_mode(False):
|
||||
dc = HfGloveDataset(
|
||||
dataset_name=name,
|
||||
split=split,
|
||||
glove_dim=embed_dim,
|
||||
max_seq_len=max_seq_len,
|
||||
sample_field_name=sample_field_name,
|
||||
label_field_name=label_field_name)
|
||||
return (dc,)
|
||||
@@ -0,0 +1,99 @@
|
||||
from typing import Any, Dict, Optional, Callable, Tuple
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
import torch
|
||||
from typing import Tuple
|
||||
|
||||
class SequentialTensorDataset(Dataset):
|
||||
"""
|
||||
Implementation of a sequential dataset.
|
||||
|
||||
pragma: skip_doc
|
||||
"""
|
||||
def __init__(self, tens: torch.Tensor, sequence_length: int):
|
||||
"""
|
||||
Args:
|
||||
tens (Tensor): Input tensor of shape [N, ...] containing the full sequence.
|
||||
sequence_length (int): Length of each input sequence window.
|
||||
"""
|
||||
if tens.dim() == 1:
|
||||
self.data_tens = torch.unsqueeze(tens, -1)
|
||||
else:
|
||||
self.data_tens = tens
|
||||
|
||||
self.sequence_length = sequence_length
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""
|
||||
Returns:
|
||||
int: Number of valid (input, target) sequence pairs.
|
||||
"""
|
||||
return self.data_tens.size(0) - self.sequence_length - 1
|
||||
|
||||
def __getitem__(self, i: int) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Args:
|
||||
i (int): Index of the sequence window.
|
||||
|
||||
Returns:
|
||||
Tuple[Tensor, Tensor]: A tuple containing:
|
||||
- x: Input sequence of shape [sequence_length, ...]
|
||||
- y: Target sequence of shape [sequence_length, ...], shifted by one step
|
||||
"""
|
||||
x = self.data_tens[i : i + self.sequence_length ]
|
||||
y = self.data_tens[i + 1 : i + self.sequence_length + 1]
|
||||
return x, y
|
||||
|
||||
|
||||
|
||||
class PtvSequentialTensorDataset:
|
||||
"""
|
||||
Creates a sequential tensor Dataset.
|
||||
|
||||
Args:
|
||||
tens (torch.Tensor): The tensor containing the data.
|
||||
seq_len (int): Sequence length.
|
||||
|
||||
category: PyTorch wrapper - Training
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
||||
"""
|
||||
Defines the required input types.
|
||||
|
||||
Returns:
|
||||
Dict[str, Any]: A dictionary of required input types.
|
||||
"""
|
||||
return {
|
||||
"required": {
|
||||
"tens": ("TENSOR", {}),
|
||||
"seq_len": ("INT", {"default":32, "min":1, "max":1e5})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES: Tuple[str] = ("PTVDATASET",)
|
||||
FUNCTION: str = "f"
|
||||
CATEGORY: str = "Training"
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kw):
|
||||
return float("NaN")
|
||||
|
||||
def f(self,
|
||||
tens: torch.Tensor,
|
||||
seq_len: int) -> Tuple:
|
||||
"""
|
||||
Creates a sequential tensor Dataset.
|
||||
|
||||
Args:
|
||||
tens (torch.Tensor): The tensor containing the data.
|
||||
seq_len (int): Sequence length.
|
||||
|
||||
Returns:
|
||||
Tuple: A tuple containing the dataset instance.
|
||||
"""
|
||||
|
||||
with torch.inference_mode(False):
|
||||
dc = SequentialTensorDataset(tens, seq_len)
|
||||
return (dc,)
|
||||
@@ -0,0 +1,84 @@
|
||||
# Return Value of RNN Model's forward()
|
||||
"""
|
||||
model = nn.RNN(...) # or nn.LSTM, nn.GRU
|
||||
out_seq, _ = model(input)
|
||||
|
||||
out_seq has shape:
|
||||
(batch, seq_len, hidden_size) if batch_first=True
|
||||
|
||||
To extract the output from the last time step of the top RNN layer:
|
||||
|
||||
Unidirectional Case:
|
||||
final_output = out_seq[:, -1, :] # shape: (batch, hidden_size)
|
||||
|
||||
Bidirectional Case:
|
||||
If the model is bidirectional, hidden_size in out_seq is already doubled.
|
||||
No need to manually concatenate forward/backward states.
|
||||
|
||||
final_output = out_seq[:, -1, :] # shape: (batch, 2 * hidden_size)
|
||||
|
||||
To match label shape:
|
||||
- If the label is 3D: shape (batch, seq_len, hidden_size), use label[:, -1, :]
|
||||
- If bidirectional: repeat the last axis
|
||||
"""
|
||||
from typing import Tuple, Union
|
||||
import torch
|
||||
|
||||
def extract_rnn_return_value_with_adjusted_label(
|
||||
inputs: torch.Tensor,
|
||||
label: torch.Tensor,
|
||||
rnn_output: torch.Tensor,
|
||||
bidirectional: bool,
|
||||
batch_first:bool = True,
|
||||
process_label:bool=True,
|
||||
return_valid_token_mean:bool=True
|
||||
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
|
||||
"""
|
||||
Extracts the final output at the last time step from the top RNN layer (out_seq[:, -1, :])
|
||||
and adjusts the label shape to match it.
|
||||
|
||||
Args:
|
||||
inputs (torch.Tensor): The input (feature) tensor.
|
||||
label (torch.Tensor): Ground truth tensor of shape (batch, seq_len, hidden_size).
|
||||
rnn_output (torch.Tensor): The output sequence from the RNN (out_seq),
|
||||
shape (batch, seq_len, hidden_size) for unidirectional,
|
||||
or (batch, seq_len, 2 * hidden_size) for bidirectional RNNs.
|
||||
bidirectional (bool): Whether the RNN is bidirectional.
|
||||
batch_first (bool): True if the first axis is the batch.
|
||||
process_label (bool): True if process label. To process only the RNN output, pass None to label and set this flag to False.
|
||||
return_valid_token_mean (bool): True to return the mean of non-zero outputs.
|
||||
|
||||
Returns:
|
||||
Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
|
||||
- The final time-step output of shape (batch, hidden_size) or (batch, 2 * hidden_size)
|
||||
- The adjusted label of matching shape
|
||||
"""
|
||||
if batch_first:
|
||||
if return_valid_token_mean:
|
||||
# Use inputs to build a valid token mask
|
||||
valid_mask = inputs.norm(p=2, dim=2) > 0 # shape: [batch, seq_len]
|
||||
lengths = valid_mask.sum(dim=1).clamp(min=1) # shape: [batch]
|
||||
masked = rnn_output * valid_mask.unsqueeze(-1) # [batch, seq_len, hidden]
|
||||
out = masked.sum(dim=1) / lengths.unsqueeze(1) # [batch, hidden]
|
||||
else:
|
||||
out = rnn_output[:, -1, :] # last time step output from top layer
|
||||
if process_label:
|
||||
label = label[:, -1, :]
|
||||
else:
|
||||
if return_valid_token_mean:
|
||||
valid_mask = inputs.norm(p=2, dim=2) > 0 # shape: [seq_len, batch]
|
||||
lengths = valid_mask.sum(dim=0).clamp(min=1) # shape: [batch]
|
||||
masked = rnn_output * valid_mask.unsqueeze(-1) # [seq_len, batch, hidden]
|
||||
out = masked.sum(dim=0) / lengths.unsqueeze(1) # [batch, hidden]
|
||||
else:
|
||||
out = rnn_output[-1, :, :] # last time step output from top layer
|
||||
if process_label:
|
||||
label = label[-1, :, :]
|
||||
|
||||
if bidirectional and process_label:
|
||||
label = torch.cat([label, label], dim=-1)
|
||||
|
||||
if process_label:
|
||||
return out, label
|
||||
else:
|
||||
return out
|
||||
@@ -1,6 +1,8 @@
|
||||
import os
|
||||
import ast
|
||||
from typing import Optional, Tuple, Union
|
||||
import random
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
@@ -137,3 +139,17 @@ def str_to_dim(dim: str) -> Optional[Union[int, Tuple[int, ...]]]:
|
||||
|
||||
return None
|
||||
|
||||
def set_seed(seed=42) -> None:
|
||||
"""
|
||||
Sets seed for random number generation.
|
||||
|
||||
Args:
|
||||
seed (int): Initial seed.
|
||||
"""
|
||||
torch.manual_seed(seed)
|
||||
np.random.seed(seed)
|
||||
random.seed(seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
Reference in New Issue
Block a user