82 lines
2.5 KiB
Python
82 lines
2.5 KiB
Python
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,)
|