dataeval 0.74.2__py3-none-any.whl → 0.75.0__py3-none-any.whl

dataeval/utils/metadata.py

@@ -1,12 +1,27 @@
+"""
+Metadata related utility functions that help organize raw metadata into :class:`Metadata` objects
+for use within `DataEval`.
+"""
+
 from __future__ import annotations
 
-__all__ = ["merge_metadata"]
+__all__ = ["Metadata", "preprocess", "merge"]
 
 import warnings
-from typing import Any, Iterable, Mapping, TypeVar, overload
+from dataclasses import dataclass
+from typing import Any, Iterable, Literal, Mapping, TypeVar, overload
 
 import numpy as np
-from numpy.typing import NDArray
+from numpy.typing import ArrayLike, NDArray
+from scipy.stats import wasserstein_distance as wd
+
+from dataeval.interop import as_numpy, to_numpy
+from dataeval.output import Output, set_metadata
+
+TNum = TypeVar("TNum", int, float)
+DISCRETE_MIN_WD = 0.054
+CONTINUOUS_MIN_SAMPLE_SIZE = 20
+
 
 T = TypeVar("T")
 
@@ -185,7 +200,7 @@ def _is_metadata_dict_of_dicts(metadata: Mapping) -> bool:
     return set(metadata[keys[0]]) == set(metadata[keys[1]])
 
 
-def merge_metadata(
+def merge(
     metadata: Iterable[Mapping[str, Any]],
     ignore_lists: bool = False,
     fully_qualified: bool = False,
@@ -222,7 +237,7 @@ def merge_metadata(
     Example
     -------
     >>> list_metadata = [{"common": 1, "target": [{"a": 1, "b": 3, "c": 5}, {"a": 2, "b": 4}], "source": "example"}]
-    >>> reorganized_metadata, image_indicies = merge_metadata(list_metadata)
+    >>> reorganized_metadata, image_indicies = merge(list_metadata)
     >>> reorganized_metadata
     {'common': [1, 1], 'a': [1, 2], 'b': [3, 4], 'source': ['example', 'example']}
     >>> image_indicies
@@ -276,3 +291,293 @@ def merge_metadata(
         output[k] = np.array(cv) if as_numpy else cv
 
     return output, image_indicies
+
+
+@dataclass(frozen=True)
+class Metadata(Output):
+    """
+    Dataclass containing binned metadata from the :func:`preprocess` function
+
+    Attributes
+    ----------
+    discrete_factor_names : list[str]
+        List containing factor names for the original data that was discrete and the binned continuous data
+    discrete_data : NDArray[np.int]
+        Array containing values for the original data that was discrete and the binned continuous data
+    continuous_factor_names : list[str]
+        List containing factor names for the original continuous data
+    continuous_data : NDArray[np.int or np.double] | None
+        Array containing values for the original continuous data or None if there was no continuous data
+    class_labels : NDArray[np.int]
+        Numerical class labels for the images/objects
+    class_names : NDArray[Any]
+        Array of unique class names (for use with plotting)
+    total_num_factors : int
+        Sum of discrete_factor_names and continuous_factor_names plus 1 for class
+    """
+
+    discrete_factor_names: list[str]
+    discrete_data: NDArray[np.int_]
+    continuous_factor_names: list[str]
+    continuous_data: NDArray[np.int_ | np.double] | None
+    class_labels: NDArray[np.int_]
+    class_names: NDArray[Any]
+    total_num_factors: int
+
+
+@set_metadata
+def preprocess(
+    raw_metadata: Iterable[Mapping[str, Any]],
+    class_labels: ArrayLike | str,
+    continuous_factor_bins: Mapping[str, int | list[tuple[TNum, TNum]]] | None = None,
+    auto_bin_method: Literal["uniform_width", "uniform_count", "clusters"] = "uniform_width",
+    exclude: Iterable[str] | None = None,
+) -> Metadata:
+    """
+    Restructures the metadata to be in the correct format for the bias functions.
+
+    This identifies whether the incoming metadata is discrete or continuous,
+    and whether the data is already binned or still needs binning.
+    It accepts a list of dictionaries containing the per image metadata and
+    automatically adjusts for multiple targets in an image.
+
+    Parameters
+    ----------
+    raw_metadata : Iterable[Mapping[str, Any]]
+        Iterable collection of metadata dictionaries to flatten and merge.
+    class_labels : ArrayLike or string
+        If arraylike, expects the labels for each image (image classification) or each object (object detection).
+        If the labels are included in the metadata dictionary, pass in the key value.
+    continuous_factor_bins : Mapping[str, int] or Mapping[str, list[tuple[TNum, TNum]]] or None, default None
+        User provided dictionary specifying how to bin the continuous metadata factors
+    auto_bin_method : "uniform_width" or "uniform_count" or "clusters", default "uniform_width"
+        Method by which the function will automatically bin continuous metadata factors. It is recommended
+        that the user provide the bins through the `continuous_factor_bins`.
+    exclude : Iterable[str] or None, default None
+        User provided collection of metadata keys to exclude when processing metadata.
+
+    Returns
+    -------
+    Metadata
+        Output class containing the binned metadata
+    """
+    # Transform metadata into single, flattened dictionary
+    metadata, image_repeats = merge(raw_metadata)
+
+    # Drop any excluded metadata keys
+    if exclude:
+        for k in list(metadata):
+            if k in exclude:
+                metadata.pop(k)
+
+    # Get the class label array in numeric form
+    class_array = as_numpy(metadata.pop(class_labels)) if isinstance(class_labels, str) else as_numpy(class_labels)
+    if class_array.ndim > 1:
+        raise ValueError(
+            f"Got class labels with {class_array.ndim}-dimensional "
+            f"shape {class_array.shape}, but expected a 1-dimensional array."
+        )
+    if not np.issubdtype(class_array.dtype, np.int_):
+        unique_classes, numerical_labels = np.unique(class_array, return_inverse=True)
+    else:
+        numerical_labels = class_array
+        unique_classes = np.unique(class_array)
+
+    # Bin according to user supplied bins
+    continuous_metadata = {}
+    discrete_metadata = {}
+    if continuous_factor_bins is not None and continuous_factor_bins != {}:
+        invalid_keys = set(continuous_factor_bins.keys()) - set(metadata.keys())
+        if invalid_keys:
+            raise KeyError(
+                f"The keys - {invalid_keys} - are present in the `continuous_factor_bins` dictionary "
+                "but are not keys in the `metadata` dictionary. Delete these keys from `continuous_factor_bins` "
+                "or add corresponding entries to the `metadata` dictionary."
+            )
+        for factor, grouping in continuous_factor_bins.items():
+            discrete_metadata[factor] = _user_defined_bin(metadata[factor], grouping)
+            continuous_metadata[factor] = metadata[factor]
+
+    # Determine category of the rest of the keys
+    remaining_keys = set(metadata.keys()) - set(continuous_metadata.keys())
+    for key in remaining_keys:
+        data = to_numpy(metadata[key])
+        if np.issubdtype(data.dtype, np.number):
+            result = _is_continuous(data, image_repeats)
+            if result:
+                continuous_metadata[key] = data
+            unique_samples, ordinal_data = np.unique(data, return_inverse=True)
+            if unique_samples.size <= np.max([20, data.size * 0.01]):
+                discrete_metadata[key] = ordinal_data
+            else:
+                warnings.warn(
+                    f"A user defined binning was not provided for {key}. "
+                    f"Using the {auto_bin_method} method to discretize the data. "
+                    "It is recommended that the user rerun and supply the desired "
+                    "bins using the continuous_factor_bins parameter.",
+                    UserWarning,
+                )
+                discrete_metadata[key] = _binning_function(data, auto_bin_method)
+        else:
+            _, discrete_metadata[key] = np.unique(data, return_inverse=True)
+
+    # splitting out the dictionaries into the keys and values
+    discrete_factor_names = list(discrete_metadata.keys())
+    discrete_data = np.stack(list(discrete_metadata.values()), axis=-1)
+    continuous_factor_names = list(continuous_metadata.keys())
+    continuous_data = np.stack(list(continuous_metadata.values()), axis=-1) if continuous_metadata else None
+    total_num_factors = len(discrete_factor_names + continuous_factor_names) + 1
+
+    return Metadata(
+        discrete_factor_names,
+        discrete_data,
+        continuous_factor_names,
+        continuous_data,
+        numerical_labels,
+        unique_classes,
+        total_num_factors,
+    )
+
+
+def _user_defined_bin(data: list[Any] | NDArray[Any], binning: int | list[tuple[TNum, TNum]]) -> NDArray[np.intp]:
+    """
+    Digitizes a list of values into a given number of bins.
+
+    Parameters
+    ----------
+    data : list | NDArray
+        The values to be digitized.
+    binning :  int | list[tuple[TNum, TNum]]
+        The number of bins for the discrete values that data will be digitized into.
+
+    Returns
+    -------
+    NDArray[np.intp]
+        The digitized values
+    """
+
+    if not np.all([np.issubdtype(type(n), np.number) for n in data]):
+        raise TypeError(
+            "Encountered a data value with non-numeric type when digitizing a factor. "
+            "Ensure all occurrences of continuous factors are numeric types."
+        )
+    if type(binning) is int:
+        _, bin_edges = np.histogram(data, bins=binning)
+        bin_edges[-1] = np.inf
+        bin_edges[0] = -np.inf
+    else:
+        bin_edges = binning
+    return np.digitize(data, bin_edges)
+
+
+def _binning_function(data: NDArray[Any], bin_method: str) -> NDArray[np.int_]:
+    """
+    Bins continuous data through either equal width bins, equal amounts in each bin, or by clusters.
+    """
+    if bin_method == "clusters":
+        # bin_edges = _binning_by_clusters(data)
+        warnings.warn(
+            "Binning by clusters is currently unavailable until changes to the clustering function go through.",
+            UserWarning,
+        )
+        bin_method = "uniform_width"
+
+    if bin_method != "clusters":
+        counts, bin_edges = np.histogram(data, bins="auto")
+        n_bins = counts.size
+        if counts[counts > 0].min() < 10:
+            for _ in range(20):
+                n_bins -= 1
+                counts, bin_edges = np.histogram(data, bins=n_bins)
+                if counts[counts > 0].min() >= 10 or n_bins < 2:
+                    break
+
+        if bin_method == "uniform_count":
+            quantiles = np.linspace(0, 100, n_bins + 1)
+            bin_edges = np.asarray(np.percentile(data, quantiles))
+
+    bin_edges[0] = -np.inf  # type: ignore # until the clusters speed up is merged
+    bin_edges[-1] = np.inf  # type: ignore # and the _binning_by_clusters can be uncommented
+    return np.digitize(data, bin_edges)  # type: ignore
+
+
+def _is_continuous(data: NDArray[np.number], image_indicies: NDArray[np.number]) -> bool:
+    """
+    Determines whether the data is continuous or discrete using the Wasserstein distance.
+
+    Given a 1D sample, we consider the intervals between adjacent points. For a continuous distribution,
+    a point is equally likely to lie anywhere in the interval bounded by its two neighbors. Furthermore,
+    we can put all "between neighbor" locations on the same scale of 0 to 1 by subtracting the smaller
+    neighbor and dividing out the length of the interval. (Duplicates are either assigned to zero or
+    ignored, depending on context). These normalized locations will be much more uniformly distributed
+    for continuous data than for discrete, and this gives us a way to distinguish them. Call this the
+    Normalized Near Neighbor distribution (NNN), defined on the interval [0,1].
+
+    The Wasserstein distance is available in scipy.stats.wasserstein_distance. We can use it to measure
+    how close the NNN is to a uniform distribution over [0,1]. We found that as long as a sample has at
+    least 20 points, and furthermore at least half as many points as there are discrete values, we can
+    reliably distinguish discrete from continuous samples by testing that the Wasserstein distance
+    measured from a uniform distribution is greater or less than 0.054, respectively.
+    """
+    # Check if the metadata is image specific
+    _, data_indicies_unsorted = np.unique(data, return_index=True)
+    if data_indicies_unsorted.size == image_indicies.size:
+        data_indicies = np.sort(data_indicies_unsorted)
+        if (data_indicies == image_indicies).all():
+            data = data[data_indicies]
+
+    # OLD METHOD
+    # uvals = np.unique(data)
+    # pct_unique = uvals.size / data.size
+    # return pct_unique < threshold
+
+    n_examples = len(data)
+
+    if n_examples < CONTINUOUS_MIN_SAMPLE_SIZE:
+        warnings.warn(
+            f"All samples look discrete with so few data points (< {CONTINUOUS_MIN_SAMPLE_SIZE})", UserWarning
+        )
+        return False
+
+    # Require at least 3 unique values before bothering with NNN
+    xu = np.unique(data, axis=None)
+    if xu.size < 3:
+        return False
+
+    Xs = np.sort(data)
+
+    X0, X1 = Xs[0:-2], Xs[2:]  # left and right neighbors
+
+    dx = np.zeros(n_examples - 2)  # no dx at end points
+    gtz = (X1 - X0) > 0  # check for dups; dx will be zero for them
+    dx[np.logical_not(gtz)] = 0.0
+
+    dx[gtz] = (Xs[1:-1] - X0)[gtz] / (X1 - X0)[gtz]  # the core idea: dx is NNN samples.
+
+    shift = wd(dx, np.linspace(0, 1, dx.size))  # how far is dx from uniform, for this feature?
+
+    return shift < DISCRETE_MIN_WD  # if NNN is close enough to uniform, consider the sample continuous.
+
+
+def get_counts(data: NDArray[np.int_], min_num_bins: int | None = None) -> NDArray[np.int_]:
+    """
+    Returns columnwise unique counts for discrete data.
+
+    Parameters
+    ----------
+    data : NDArray
+        Array containing integer values for metadata factors
+    min_num_bins : int | None, default None
+        Minimum number of bins for bincount, helps force consistency across runs
+
+    Returns
+    -------
+    NDArray[np.int_]
+        Bin counts per column of data.
+    """
+    max_value = data.max() + 1 if min_num_bins is None else min_num_bins
+    cnt_array = np.zeros((max_value, data.shape[1]), dtype=np.int_)
+    for idx in range(data.shape[1]):
+        cnt_array[:, idx] = np.bincount(data[:, idx], minlength=max_value)
+
+    return cnt_array

dataeval/{metrics/bias/metadata_utils.py → utils/plot.py}

@@ -3,10 +3,9 @@ from __future__ import annotations
 __all__ = []
 
 import contextlib
-from typing import Any
 
 import numpy as np
-from numpy.typing import ArrayLike, NDArray
+from numpy.typing import ArrayLike
 
 from dataeval.interop import to_numpy
 
@@ -14,30 +13,6 @@ with contextlib.suppress(ImportError):
     from matplotlib.figure import Figure
 
 
-def get_counts(data: NDArray[np.int_], min_num_bins: int | None = None) -> NDArray[np.int_]:
-    """
-    Returns columnwise unique counts for discrete data.
-
-    Parameters
-    ----------
-    data : NDArray
-        Array containing integer values for metadata factors
-    min_num_bins : int | None, default None
-        Minimum number of bins for bincount, helps force consistency across runs
-
-    Returns
-    -------
-    NDArray[np.int_]
-        Bin counts per column of data.
-    """
-    max_value = data.max() + 1 if min_num_bins is None else min_num_bins
-    cnt_array = np.zeros((max_value, data.shape[1]), dtype=np.int_)
-    for idx in range(data.shape[1]):
-        cnt_array[:, idx] = np.bincount(data[:, idx], minlength=max_value)
-
-    return cnt_array
-
-
 def heatmap(
     data: ArrayLike,
     row_labels: list[str] | ArrayLike,
@@ -149,81 +124,3 @@ def format_text(*args: str) -> str:
     """
     x = args[0]
     return f"{x:.2f}".replace("0.00", "0").replace("0.", ".").replace("nan", "")
-
-
-def diversity_bar_plot(labels: NDArray[Any], bar_heights: NDArray[Any]) -> Figure:
-    """
-    Plots a formatted bar plot
-
-    Parameters
-    ----------
-    labels : NDArray
-        Array containing the labels for each bar
-    bar_heights : NDArray
-        Array containing the values for each bar
-
-    Returns
-    -------
-    matplotlib.figure.Figure
-        Bar plot figure
-    """
-    import matplotlib.pyplot as plt
-
-    fig, ax = plt.subplots(figsize=(10, 10))
-
-    ax.bar(labels, bar_heights)
-    ax.set_xlabel("Factors")
-
-    plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")
-
-    fig.tight_layout()
-    return fig
-
-
-def coverage_plot(images: NDArray[Any], num_images: int) -> Figure:
-    """
-    Creates a single plot of all of the provided images
-
-    Parameters
-    ----------
-    images : NDArray
-        Array containing only the desired images to plot
-
-    Returns
-    -------
-    matplotlib.figure.Figure
-        Plot of all provided images
-    """
-    import matplotlib.pyplot as plt
-
-    num_images = min(num_images, len(images))
-
-    if images.ndim == 4:
-        images = np.moveaxis(images, 1, -1)
-    elif images.ndim == 3:
-        images = np.repeat(images[:, :, :, np.newaxis], 3, axis=-1)
-    else:
-        raise ValueError(
-            f"Expected a (N,C,H,W) or a (N, H, W) set of images, but got a {images.ndim}-dimensional set of images."
-        )
-
-    rows = int(np.ceil(num_images / 3))
-    fig, axs = plt.subplots(rows, 3, figsize=(9, 3 * rows))
-
-    if rows == 1:
-        for j in range(3):
-            if j >= len(images):
-                continue
-            axs[j].imshow(images[j])
-            axs[j].axis("off")
-    else:
-        for i in range(rows):
-            for j in range(3):
-                i_j = i * 3 + j
-                if i_j >= len(images):
-                    continue
-                axs[i, j].imshow(images[i_j])
-                axs[i, j].axis("off")
-
-    fig.tight_layout()
-    return fig

dataeval/utils/torch/__init__.py

@@ -5,21 +5,6 @@ While these metrics can take in custom models, DataEval provides utility classes
 to create a seamless integration between custom models and DataEval's metrics.
 """
 
-from dataeval import _IS_TORCH_AVAILABLE, _IS_TORCHVISION_AVAILABLE
+__all__ = ["models", "trainer"]
 
-__all__ = []
-
-if _IS_TORCH_AVAILABLE:
-    from dataeval.utils.torch import models, trainer
-    from dataeval.utils.torch.utils import read_dataset
-
-    __all__ += ["read_dataset", "models", "trainer"]
-
-if _IS_TORCHVISION_AVAILABLE:
-    from dataeval.utils.torch import datasets
-
-    __all__ += ["datasets"]
-
-
-del _IS_TORCH_AVAILABLE
-del _IS_TORCHVISION_AVAILABLE
+from dataeval.utils.torch import models, trainer

dataeval/utils/torch/gmm.py

@@ -1,7 +1,5 @@
 """
-Adapted for Pytorch from:
-
-Source code derived from Alibi-Detect 0.11.4
+Adapted for Pytorch from Alibi-Detect 0.11.4
 https://github.com/SeldonIO/alibi-detect/tree/v0.11.4
 
 Original code Copyright (c) 2023 Seldon Technologies Ltd
@@ -10,13 +8,38 @@ Licensed under Apache Software License (Apache 2.0)
 
 from __future__ import annotations
 
+__all__ = []
+
+from typing import NamedTuple, TypeVar
+
 import numpy as np
 import torch
 
-from dataeval.utils.gmm import GaussianMixtureModelParams
+TGMMData = TypeVar("TGMMData")
+
+
+class GaussianMixtureModelParams(NamedTuple):
+    """
+    phi : torch.Tensor
+        Mixture component distribution weights.
+    mu : torch.Tensor
+        Mixture means.
+    cov : torch.Tensor
+        Mixture covariance.
+    L : torch.Tensor
+        Cholesky decomposition of `cov`.
+    log_det_cov : torch.Tensor
+        Log of the determinant of `cov`.
+    """
+
+    phi: torch.Tensor
+    mu: torch.Tensor
+    cov: torch.Tensor
+    L: torch.Tensor
+    log_det_cov: torch.Tensor
 
 
-def gmm_params(z: torch.Tensor, gamma: torch.Tensor) -> GaussianMixtureModelParams[torch.Tensor]:
+def gmm_params(z: torch.Tensor, gamma: torch.Tensor) -> GaussianMixtureModelParams:
     """
     Compute parameters of Gaussian Mixture Model.
 
@@ -58,7 +81,7 @@ def gmm_params(z: torch.Tensor, gamma: torch.Tensor) -> GaussianMixtureModelPara
 
 def gmm_energy(
     z: torch.Tensor,
-    params: GaussianMixtureModelParams[torch.Tensor],
+    params: GaussianMixtureModelParams,
     return_mean: bool = True,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """

dataeval/utils/torch/{utils.py → internal.py}

@@ -1,70 +1,15 @@
 from __future__ import annotations
 
-__all__ = ["read_dataset"]
+__all__ = []
 
-from collections import defaultdict
 from functools import partial
 from typing import Any, Callable
 
 import numpy as np
 import torch
 from numpy.typing import NDArray
-from torch.utils.data import Dataset
-
-
-def read_dataset(dataset: Dataset[Any]) -> list[list[Any]]:
-    """
-    Extract information from a dataset at each index into individual lists of each information position
-
-    Parameters
-    ----------
-    dataset : torch.utils.data.Dataset
-        Input dataset
-
-    Returns
-    -------
-    List[List[Any]]
-        All objects in individual lists based on return position from dataset
-
-    Warning
-    -------
-    No type checking is done between lists or data inside lists
-
-    See Also
-    --------
-    torch.utils.data.Dataset
-
-    Examples
-    --------
-    >>> import numpy as np
-    >>> data = np.ones((10, 1, 3, 3))
-    >>> labels = np.ones((10,))
-    >>> class ICDataset:
-    ...     def __init__(self, data, labels):
-    ...         self.data = data
-    ...         self.labels = labels
-    ...
-    ...     def __getitem__(self, idx):
-    ...         return self.data[idx], self.labels[idx]
-
-    >>> ds = ICDataset(data, labels)
-
-    >>> result = read_dataset(ds)
-    >>> len(result)  # images and labels
-    2
-    >>> np.asarray(result[0]).shape  # images
-    (10, 1, 3, 3)
-    >>> np.asarray(result[1]).shape  # labels
-    (10,)
-    """
-
-    ddict: dict[int, list[Any]] = defaultdict(list[Any])
-
-    for data in dataset:
-        for i, d in enumerate(data if isinstance(data, tuple) else (data,)):
-            ddict[i].append(d)
-
-    return list(ddict.values())
+from torch.utils.data import DataLoader, TensorDataset
+from tqdm import tqdm
 
 
 def get_device(device: str | torch.device | None = None) -> torch.device:
@@ -167,3 +112,82 @@ def predict_batch(
         tuple(concat(p) for p in preds) if isinstance(preds, tuple) else concat(preds)  # type: ignore
     )
     return out
+
+
+def trainer(
+    model: torch.nn.Module,
+    x_train: NDArray[Any],
+    y_train: NDArray[Any] | None,
+    loss_fn: Callable[..., torch.Tensor | torch.nn.Module] | None,
+    optimizer: torch.optim.Optimizer | None,
+    preprocess_fn: Callable[[torch.Tensor], torch.Tensor] | None,
+    epochs: int,
+    batch_size: int,
+    device: torch.device,
+    verbose: bool,
+) -> None:
+    """
+    Train Pytorch model.
+
+    Parameters
+    ----------
+    model
+        Model to train.
+    loss_fn
+        Loss function used for training.
+    x_train
+        Training data.
+    y_train
+        Training labels.
+    optimizer
+        Optimizer used for training.
+    preprocess_fn
+        Preprocessing function applied to each training batch.
+    epochs
+        Number of training epochs.
+    reg_loss_fn
+        Allows an additional regularisation term to be defined as reg_loss_fn(model)
+    batch_size
+        Batch size used for training.
+    buffer_size
+        Maximum number of elements that will be buffered when prefetching.
+    verbose
+        Whether to print training progress.
+    """
+    if optimizer is None:
+        optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+
+    if y_train is None:
+        dataset = TensorDataset(torch.from_numpy(x_train).to(torch.float32))
+
+    else:
+        dataset = TensorDataset(
+            torch.from_numpy(x_train).to(torch.float32), torch.from_numpy(y_train).to(torch.float32)
+        )
+
+    loader = DataLoader(dataset=dataset)
+
+    model = model.to(device)
+
+    # iterate over epochs
+    loss = torch.nan
+    disable_tqdm = not verbose
+    for epoch in (pbar := tqdm(range(epochs), disable=disable_tqdm)):
+        epoch_loss = loss
+        for step, data in enumerate(loader):
+            if step % 250 == 0:
+                pbar.set_description(f"Epoch: {epoch} ({epoch_loss:.3f}), loss: {loss:.3f}")
+
+            x, y = [d.to(device) for d in data] if len(data) > 1 else (data[0].to(device), None)
+
+            if isinstance(preprocess_fn, Callable):
+                x = preprocess_fn(x)
+
+            y_hat = model(x)
+            y = x if y is None else y
+
+            loss = loss_fn(y, y_hat)  # type: ignore
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()