nodes/HowToSD-ComfyUI-Data-Analysis
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added experimental PyTorch model training.

Browse Source
This commit is contained in:
HowToSD
2025-03-03 18:39:52 -08:00
parent 62e4f6efb8
commit 11d3c9490c
40 changed files with 1649 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/pytorch_wrapper/pt_data_loader.py
+81
View File
@@ -0,0 +1,81 @@
import inspect
from typing import Any, Dict
import ast
import torch
from torch.utils.data import DataLoader
class PtDataLoader:
"""
Loads data from a dataset node and creates a PyTorch DataLoader.
**dataset**
Connect a PtvDataset node that provides the dataset.
**batch size**
Specifies the number of samples per batch.
**shuffle**
If True, shuffles the dataset before each epoch.
category: PyTorch wrapper - Training
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types required for data loading.
Returns:
Dict[str, Any]: A dictionary specifying required and optional input types.
"""
return {
"required": {
"dataset": ("PTVDATASET", {}),
"batch_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
"shuffle": ("BOOLEAN", {"default": False}),
},
"optional": {
"parameters": ("STRING", {"default": "", "multiline": True})
}
}
RETURN_TYPES: tuple = ("PTDATALOADER",)
FUNCTION: str = "f"
CATEGORY: str = "Training"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self,
dataset: Any, # One of the classes defined in torchvision.datasets
batch_size: int,
shuffle: bool,
parameters: str) -> tuple:
"""
Creates a PyTorch DataLoader for the given dataset.
Args:
dataset (torchvision.datasets.{MNIST, ...}): An instance of a dataset from torchvision.datasets.
batch_size (int): The number of samples per batch.
shuffle (bool): If True, shuffles the dataset before each epoch.
parameters (str): Additional parameters in Python dictionary format.
Returns:
tuple: A tuple containing the created DataLoader instance.
"""
with torch.inference_mode(False):
param_dict = ast.literal_eval(parameters)
# Get valid parameters for the DataLoader class
valid_params = inspect.signature(DataLoader).parameters
# Filter only parameters that match the dataset's constructor
filtered_params = {k: v for k, v in param_dict.items() if k in valid_params}
dl = DataLoader(dataset=dataset,
batch_size=batch_size,
shuffle=shuffle,
**filtered_params)
return (dl,)
modules/pytorch_wrapper/pt_evaluate_classification_model.py
+71
View File
@@ -0,0 +1,71 @@
from typing import Any, Dict, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
class PtEvaluateClassificationModel:
"""
Performs inference on test data and computes evaluation metrics.
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": {
"model": ("PTMODEL", {}),
"data_loader": ("PTDATALOADER", {}),
"use_gpu": ("BOOLEAN", {"default":False})
}
}
RETURN_NAMES: Tuple[str, ...] = ("accuracy", "precision", "recall", "f1", "num_samples")
RETURN_TYPES: Tuple[str, ...] = ("FLOAT", "FLOAT", "FLOAT", "FLOAT", "INT")
FUNCTION: str = "f"
CATEGORY: str = "Training"
def f(self, model: nn.Module, data_loader: torch.utils.data.DataLoader, use_gpu: bool) -> Tuple[float, float, float, float, int]:
"""
Performs inference on test data and computes evaluation metrics.
Args:
model (torch.nn.Module): The PyTorch model to evaluate.
data_loader (torch.utils.data.DataLoader): DataLoader for the test dataset.
use_gpu (bool): Whether to use GPU for inference.
Returns:
Tuple[float, float, float, float, int]: A tuple containing accuracy, precision, recall,
F1-score, and the number of samples.
"""
with torch.inference_mode():
if use_gpu:
model.to("cuda")
model.eval()
y_hat = []
y = []
for images, labels in data_loader:
if use_gpu:
images = images.to("cuda")
labels = labels.to("cuda")
outputs = model(images)
preds = outputs.argmax(dim=1)
y_hat.extend(preds.cpu().numpy())
y.extend(labels.cpu().numpy())
accuracy = accuracy_score(y, y_hat)
precision = precision_score(y, y_hat, average='macro')
recall = recall_score(y, y_hat, average='macro')
f1 = f1_score(y, y_hat, average='macro')
num_samples = len(y)
return (accuracy, precision, recall, f1, num_samples)
modules/pytorch_wrapper/pt_load_model.py
+41
View File
@@ -0,0 +1,41 @@
import os
from typing import Any, Dict
import torch
from .utils import get_model_full_path
class PtLoadModel:
"""
A wrapper class for saving a PyTorch model.
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": {
"model": ("PTMODEL", {}),
"model_path": ("STRING", {"default": "model.pkl"})
}
}
RETURN_NAMES: tuple = ("model",)
RETURN_TYPES: tuple = ("PTMODEL",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self, model:Any, model_path: str):
model_full_path = get_model_full_path(model_path)
sd = torch.load(model_full_path)
model.load_state_dict(sd)
return (model,)
modules/pytorch_wrapper/pt_save_model.py
+36
View File
@@ -0,0 +1,36 @@
import os
from typing import Any, Dict
import torch
from .utils import get_model_full_path
class PtSaveModel:
"""
A wrapper class for saving a PyTorch model.
category: IO
"""
@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": {
"model": ("PTMODEL", {}),
"model_path": ("STRING", {"default": "model.pkl"})
}
}
RETURN_NAMES: tuple = ("model full path",)
RETURN_TYPES: tuple = ("STRING",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, model:Any, model_path: str):
model_full_path = get_model_full_path(model_path)
torch.save(model.state_dict(), model_full_path)
return (os.path.realpath(model_full_path),)
modules/pytorch_wrapper/pt_train_classification_model.py
+111
View File
@@ -0,0 +1,111 @@
from typing import Any, Dict, Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
class PtTrainClassificationModel:
"""
Trains a classification model using a given dataset, optimizer, and number of epochs.
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", {}),
"epochs": ("INT", {"default":1, "min":1, "max":1e6}),
"use_gpu": ("BOOLEAN", {"default":False})
},
"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,
epochs: int,
use_gpu: 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.
epochs (int): Number of training epochs.
use_gpu (bool): Whether to use GPU for training.
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.
"""
with torch.inference_mode(False):
train_loss = list()
val_loss = list()
loss_function = nn.CrossEntropyLoss(reduction="sum")
# Train the model
for epoch in range(epochs):
if use_gpu:
model.to("cuda")
model.train()
total_loss = 0
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()
total_loss += loss.item()
epoch_loss = total_loss / len(train_loader)
print(f"Epoch (train) {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}")
train_loss.append(epoch_loss)
if val_loader:
pass
with torch.no_grad():
model.eval()
total_loss = 0
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)
total_loss += loss.item()
epoch_loss = total_loss / len(val_loader)
print(f"Epoch (val) {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}")
val_loss.append(epoch_loss)
return (model,
torch.Tensor(np.asarray(train_loss)),
torch.Tensor(np.asarray(val_loss)))
modules/pytorch_wrapper/ptn_linear_model.py
+86
View File
@@ -0,0 +1,86 @@
from typing import Any, Dict, Tuple
import ast
import torch
import torch.nn as nn
from .ptn_linear_model_def import DenseModel
class PtnLinearModel:
"""
A linear model consisting of dense layers.
### Fields:
- `dim_list`: A comma-separated string representing layer dimensions.
It should contain the input dimension for the first layer and the output dimensions
for all layers.
Example for a two-layer network:
- Input data: 784
- Layer 1: 784 → 128
- Layer 2: 128 → 10
- Specify `"784, 128, 10"`
- `bias_list`: A comma-separated string representing bias usage (True/False) per layer.
- `num_layers`: The number of layers in the model.
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": {
"dim_list": ("STRING", {"default": "784,10"}),
"bias_list": ("STRING", {"default": "{True}"}),
"num_layers": ("INT", {"default": 1, "min": 1, "max": 2000}),
}
}
RETURN_TYPES: Tuple[str] = ("PTMODEL",)
FUNCTION: str = "f"
CATEGORY: str = "Training"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self, dim_list: str, bias_list: str, num_layers: int) -> Tuple[nn.Module]:
"""
Constructs a dense neural network model.
Args:
dim_list (str): A comma-separated string representing layer dimensions.
bias_list (str): A comma-separated string representing bias usage (True/False) per layer.
num_layers (int): The number of layers in the model.
Returns:
Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
Raises:
ValueError: If `dim_list` length does not match `num_layers + 1`.
ValueError: If `bias_list` length does not match `num_layers`.
"""
dims = ast.literal_eval(dim_list)
biases = ast.literal_eval(bias_list)
if len(dims) - 1 != num_layers:
raise ValueError(
"dim_list and num_layers do not match. dim_list should contain "
"the input dimension for the first layer and output dimensions for all layers. "
"This list should be one element longer than the number of layers."
)
if len(biases) != num_layers:
raise ValueError("bias_list and num_layers do not match.")
with torch.inference_mode(False):
model = DenseModel(dims, biases, num_layers)
return (model,)
modules/pytorch_wrapper/ptn_linear_model_def.py
+99
View File
@@ -0,0 +1,99 @@
from typing import List, Dict, Any
import torch
import torch.nn as nn
class DenseSingleLayer(nn.Module):
"""A single dense layer followed by a ReLU activation.
Args:
in_features (int): Number of input features.
out_features (int): Number of output features.
bias (bool): If True, adds a learnable bias to the layer.
pragma: skip_doc
"""
@classmethod # TODO: Dummy method to suppress ComfyUI warning
def INPUT_TYPES(cls) -> Dict[str, Any]:
return {
}
def __init__(self, in_features: int, out_features: int, bias: bool) -> None:
super().__init__()
self.fc = nn.Linear(in_features, out_features, bias)
self.activation = nn.ReLU()
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
"""Forward pass of the dense layer with ReLU activation.
Args:
inputs (torch.Tensor): Input tensor of shape `(batch_size, in_features)`.
Returns:
torch.Tensor: Output tensor of shape `(batch_size, out_features)`.
"""
x = self.fc(inputs)
x = self.activation(x)
return x
class DenseModel(nn.Module):
"""A multi-layer dense neural network with ReLU activation after each dense layer, except for the final layer.
Args:
dim_list (List[int]): List of feature dimensions. Should be of length `num_layers + 1`.
bias_list (List[bool]): List indicating whether each layer should have a bias.
num_layers (int): Number of layers in the model.
Example:
For a three-layer network:
- Layer 1: input -> output1
- Layer 2: output1 -> output2
- Layer 3: output2 -> output3
`dim_list = [input_dim, output1_dim, output2_dim, output3_dim]`
pragma: skip_doc
"""
@classmethod # TODO: Dummy method to suppress ComfyUI warning
def INPUT_TYPES(cls) -> Dict[str, Any]:
return {
}
def __init__(self, dim_list: List[int], bias_list: List[bool], num_layers: int) -> None:
super().__init__()
self.no_activation_layers = nn.ModuleList([
DenseSingleLayer(dim_list[i], dim_list[i + 1], bias_list[i])
for i in range(num_layers - 1)
])
self.dim_list = dim_list
self.bias_list = bias_list
self.num_layers = num_layers
self.final_dense = nn.Linear(
dim_list[num_layers - 1], # in_features
dim_list[num_layers], # out_features
bias_list[num_layers - 1]
)
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
"""Forward pass of the dense model.
Args:
inputs (torch.Tensor): Input tensor of shape `(batch_size, dim_list[0])`.
Returns:
torch.Tensor: Output tensor of shape `(batch_size, dim_list[-1])`.
"""
x = inputs.view(-1, self.dim_list[0]) # Ensure correct input shape
for i in range(self.num_layers - 1):
x = self.no_activation_layers[i](x)
x = self.final_dense(x)
return x
modules/pytorch_wrapper/pto_adam.py
+46
View File
@@ -0,0 +1,46 @@
from typing import Any, Dict
import torch
import torch.nn as nn
import torch.nn.functional as F
class PtoAdam:
"""
Converts PyTorch tensor to NumPy ndarray.
category: PyTorch wrapper - Tensor data conversion
"""
@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": {
"model": ("PTMODEL", {}),
"learning_rate": ("FLOAT", {"default":0.001, "min":1e-10, "max":1, "step":0.0000000001}),
"beta1": ("FLOAT", {"default":0.9, "min":1e-10, "max":1, "step":0.0000000001}),
"beta2": ("FLOAT", {"default":0.999, "min":1e-10, "max":1, "step":0.0000000001}),
}
}
RETURN_TYPES: tuple = ("PTOPTIMIZER",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, model, learning_rate, beta1, beta2) -> tuple:
"""
Converts PyTorch tensor to NumPy ndarray.
Args:
tens (torch.Tensor): PyTorch Tensor
Returns:
tuple: A tuple containing the ndarray.
"""
with torch.inference_mode(False):
opt = torch.optim.Adam(model.parameters(), lr=learning_rate, betas=(beta1, beta2))
return(opt,)
modules/pytorch_wrapper/ptv_dataset.py
+124
View File
@@ -0,0 +1,124 @@
import inspect
from typing import Any, Dict, Optional, Callable, Tuple
import torch
import torchvision.datasets as datasets
import ast
from .utils import get_dataset_full_path
class PtvDataset:
"""
A Torchvision Dataset class wrapper.
Parameters
**name**
Specify a dataset class name such as MNIST, FashionMNIST in string.
This will be converted to a dataset class internally.
**download**
Set to True if you want to download.
**root**
Specify the root directory of the downloaded dataset relative to the
dataset directory of this extension, or specify the absolute path.
If you leave as a blank, the dataset will be downloaded under the `datasets` directory
of this extension.
**transform**
Data transforms e.g. transform to PyTorch tensor, normalize image.
Plug in the PTV Transforms node that contains Torchvision Transforms functionality.
**parameters**
Specify other parameters in Python dict format.
For example, if you want to set Train parameter to True to download the train set, specify:
```
{"train": True}
```
or if you want to download the test set, specify:
```
{"train": False}
```
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}),
"download": ("BOOLEAN", {"default": True}),
"root": ("STRING", {"default": "", "multiline": False}),
},
"optional": {
"transform": ("PTVTRANSFORM", {}),
"parameters": ("STRING", {"default": "", "multiline": True})
}
}
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,
download: bool,
root: str,
parameters: str,
transform: Optional[Callable] = None) -> Tuple:
"""
Loads a dataset from Torchvision with specified parameters.
Args:
name (str): Name of the dataset class (e.g., MNIST, FashionMNIST).
download (bool): Whether to download the dataset.
root (str): Root directory for storing the dataset.
parameters (str): Additional dataset parameters in stringified dictionary format.
transform (Optional[Callable]): Transformations to apply to the dataset.
Returns:
Tuple: A tuple containing the dataset instance.
"""
with torch.inference_mode(False):
param_dict = ast.literal_eval(parameters)
# Converts the datasets class name to dataset class itself
try:
dataset_class = getattr(datasets, name)
except AttributeError:
raise ValueError(f"Dataset '{name}' is not available in torchvision.datasets")
# Get valid parameters for the dataset class
valid_params = inspect.signature(dataset_class).parameters
# Filter only parameters that match the dataset's constructor
filtered_params = {k: v for k, v in param_dict.items() if k in valid_params}
if root:
dataset_path = get_dataset_full_path(root)
else: # Use default directory
dataset_path = get_dataset_full_path("")
if transform:
dc = dataset_class(download=download,
root=dataset_path,
transform=transform,
**filtered_params)
else:
dc = dataset_class(download=download,
root=dataset_path,
**filtered_params)
return (dc,)
modules/pytorch_wrapper/ptv_dataset_loader.py
+148
View File
@@ -0,0 +1,148 @@
import ast
import inspect
from typing import Any, Dict, Tuple
import torch
import torchvision.datasets as datasets
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from .utils import get_dataset_full_path
class PtvDatasetLoader:
"""
A node to combine the dataset and data loader into a single node.
Transform is internally called to convert input data to PyTorch tensors.
Parameters
**name**
Specify the dataset class name such as MNIST, FashionMNIST in string.
This will be converted to a dataset class internally.
**download**
Set to True if you want to download.
**root**
Specify the root directory of the downloaded dataset relative to the
dataset directory of this extension, or specify the absolute path.
If you leave as a blank, the dataset will be downloaded under the `datasets` directory
of this extension.
**batch_size**
The number of samples per batch.
**shuffle**
If True, shuffles the dataset before each epoch.
**dataset parameters**
Specify other parameters in Python dict format.
For example, if you want to set Train parameter to True to download the train set, specify:
```
{"train": True}
```
or if you want to download the test set, specify:
```
{"train": False}
```
**load parameters**
Additional parameters in Python dictionary format.
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}),
"download": ("BOOLEAN", {"default": True}),
"root": ("STRING", {"default": "", "multiline": False}),
"batch_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
"shuffle": ("BOOLEAN", {"default": False}),
},
"optional": {
"dataset_parameters": ("STRING", {"default": '{"train":True}', "multiline": True}),
"load_parameters": ("STRING", {"default": '{"num_workers:1}', "multiline": True})
}
}
RETURN_TYPES: Tuple[str] = ("PTDATALOADER",)
FUNCTION: str = "f"
CATEGORY: str = "Training"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self,
name: str,
download: bool,
root: str,
batch_size: int,
shuffle: bool,
dataset_parameters: str,
load_parameters: str) -> Tuple:
"""
Loads a dataset from Torchvision with specified parameters.
Args:
name (str): Name of the dataset class (e.g., MNIST, FashionMNIST).
download (bool): Whether to download the dataset.
root (str): Root directory for storing the dataset.
batch_size (int): The number of samples per batch.
shuffle (bool): If True, shuffles the dataset before each epoch.
dataset_parameters (str): Additional dataset parameters in stringified dictionary format.
load_parameters (str): Additional parameters for data load in Python dictionary format.
Returns:
Tuple: A tuple containing the dataset instance.
"""
with torch.inference_mode(False):
dataset_param_dict = ast.literal_eval(dataset_parameters)
# Converts the datasets class name to dataset class itself
try:
dataset_class = getattr(datasets, name)
except AttributeError:
raise ValueError(f"Dataset '{name}' is not available in torchvision.datasets")
# Get valid parameters for the dataset class
valid_params = inspect.signature(dataset_class).parameters
# Filter only parameters that match the dataset's constructor
filtered_params = {k: v for k, v in dataset_param_dict.items() if k in valid_params}
if root:
dataset_path = get_dataset_full_path(root)
else: # Use default directory
dataset_path = get_dataset_full_path("")
transform = transforms.Compose([transforms.ToTensor()])
dc = dataset_class(download=download,
root=dataset_path,
transform=transform,
**filtered_params)
load_param_dict = ast.literal_eval(load_parameters)
# Get valid parameters for the DataLoader class
valid_params = inspect.signature(DataLoader).parameters
# Filter only parameters that match the dataset's constructor
filtered_params = {k: v for k, v in load_param_dict.items() if k in valid_params}
dl = DataLoader(dataset=dc,
batch_size=batch_size,
shuffle=shuffle,
**filtered_params)
return (dl,)
modules/pytorch_wrapper/ptv_transforms_to_tensor.py
+37
View File
@@ -0,0 +1,37 @@
from typing import Any, Dict, Optional, Callable
import torchvision.transforms as transforms
class PtvTransformsToTensor:
"""
Transforms elements of dataset to PyTorch tensors.
category: PyTorch wrapper - Training
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: An empty dictionary.
"""
return {
}
RETURN_TYPES: tuple = ("PTVTRANSFORM",)
FUNCTION: str = "f"
CATEGORY: str = "Training"
def f(self) -> tuple:
"""
Transforms elements of dataset to PyTorch tensors.
Returns:
tuple: A tuple containing the transform callable.
"""
c = transforms.Compose([transforms.ToTensor()])
return(c,)
modules/pytorch_wrapper/utils.py
+54
View File
@@ -1,3 +1,4 @@
import os
import ast
from typing import Optional, Tuple, Union
import torch
@@ -16,6 +17,59 @@ DTYPE_MAPPING = {
}
def get_resource_full_path(resource_name: str, resource_path: str) -> str:
"""
Resolves the absolute path to a resource file within the "resources" directory and
ensures that the directory exists.
Args:
resource_path (str): The path to the resource file. If the path is a relative path,
the path is considered relative to the `resource` directory of this extension.
Returns:
str: The absolute path to the resource file.
"""
if os.path.isabs(resource_path):
resource_full_path = resource_path
else:
resource_full_path = os.path.realpath(
os.path.join(os.path.dirname(__file__), "..", "..", resource_name, resource_path))
resource_dir = os.path.dirname(resource_full_path)
if os.path.exists(resource_dir) is False:
os.makedirs(resource_dir)
return resource_full_path
def get_model_full_path(model_path: str) -> str:
"""
Resolves the absolute path to a model file within the "models" directory and
ensures that the directory exists.
Args:
model_path (str): The path to the model file. If the path is a relative path,
the path is considered relative to the `models` directory of this extension.
Returns:
str: The absolute path to the model file.
"""
return get_resource_full_path("models", model_path)
def get_dataset_full_path(dataset_path: str) -> str:
"""
Resolves the absolute path to the datasete within the "datasets" directory and
ensures that the directory exists.
Args:
dataset_path (str): The path to the dataset file. If the path is a relative path,
the path is considered relative to the `datasets` directory of this extension.
Returns:
str: The absolute path to the dataset file.
"""
return get_resource_full_path("datasets", dataset_path)
def str_to_dim(dim: str) -> Optional[Union[int, Tuple[int, ...]]]:
"""
Converts a string representation of a dimension into an integer or a tuple of integers.