dataeval 0.73.0__py3-none-any.whl → 0.74.0__py3-none-any.whl

dataeval/utils/split_dataset.py

@@ -144,7 +144,7 @@ def check_groups(group_ids: NDArray[np.int_], num_partitions: int) -> bool:
     ----------
     group_ids : np.ndarray
         Identifies the group to which a sample at the same index belongs.
-    num_partitions: int
+    num_partitions : int
         How many total (train, val) folds will be generated (+1 if also specifying a test fold).
 
     Warns
@@ -242,12 +242,12 @@ def get_group_ids(metadata: dict[str, Any], group_names: list[str], num_samples:
 
     Returns
     -------
-    group_ids: np.ndarray
+    group_ids : np.ndarray
         group identifiers from metadata
     """
     features2group = {k: np.array(v) for k, v in metadata.items() if k in group_names}
     if not features2group:
-        return np.zeros(num_samples, dtype=int)
+        return np.zeros(num_samples, dtype=np.int_)
     for name, feature in features2group.items():
         if len(feature) != num_samples:
             raise IndexError(f"""Feature length does not match number of labels. 
@@ -300,7 +300,13 @@ def make_splits(
         splits = splitter.split(index, labels)
     for train_idx, eval_idx in splits:
         test_ratio = len(eval_idx) / index.shape[0]
-        split_defs.append({"train": train_idx.astype(int), "eval": eval_idx.astype(int), "eval_frac": test_ratio})
+        split_defs.append(
+            {
+                "train": train_idx.astype(np.int_),
+                "eval": eval_idx.astype(np.int_),
+                "eval_frac": test_ratio,
+            }
+        )
     return split_defs
 
 
@@ -318,9 +324,9 @@ def find_best_split(
     split_defs : list[dict]
         List of dictionaries, which specifying train index, validation index, and the ratio of
         validation to all data.
-    stratified: bool
+    stratified : bool
         If True, maintain dataset class balance within each train/val split
-    eval_frac: float
+    eval_frac : float
         Desired fraction of the dataset sequestered for evaluation
 
     Returns

dataeval/utils/tensorflow/_internal/gmm.py

@@ -8,10 +8,11 @@ Licensed under Apache Software License (Apache 2.0)
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, NamedTuple
+from typing import TYPE_CHECKING
 
 import numpy as np
 
+from dataeval.utils.gmm import GaussianMixtureModelParams
 from dataeval.utils.lazy import lazyload
 
 if TYPE_CHECKING:
@@ -20,28 +21,7 @@ else:
     tf = lazyload("tensorflow")
 
 
-class GaussianMixtureModelParams(NamedTuple):
-    """
-    phi : tf.Tensor
-        Mixture component distribution weights.
-    mu : tf.Tensor
-        Mixture means.
-    cov : tf.Tensor
-        Mixture covariance.
-    L : tf.Tensor
-        Cholesky decomposition of `cov`.
-    log_det_cov : tf.Tensor
-        Log of the determinant of `cov`.
-    """
-
-    phi: tf.Tensor
-    mu: tf.Tensor
-    cov: tf.Tensor
-    L: tf.Tensor
-    log_det_cov: tf.Tensor
-
-
-def gmm_params(z: tf.Tensor, gamma: tf.Tensor) -> GaussianMixtureModelParams:
+def gmm_params(z: tf.Tensor, gamma: tf.Tensor) -> GaussianMixtureModelParams[tf.Tensor]:
     """
     Compute parameters of Gaussian Mixture Model.
 
@@ -81,7 +61,7 @@ def gmm_params(z: tf.Tensor, gamma: tf.Tensor) -> GaussianMixtureModelParams:
 
 def gmm_energy(
     z: tf.Tensor,
-    params: GaussianMixtureModelParams,
+    params: GaussianMixtureModelParams[tf.Tensor],
     return_mean: bool = True,
 ) -> tuple[tf.Tensor, tf.Tensor]:
     """

dataeval/utils/torch/datasets.py

@@ -206,7 +206,7 @@ class MNIST(Dataset[tuple[NDArray[np.float64], int]]):
             Option to select specific classes from dataset.
         balance : bool, default True
             If True, returns equal number of samples for each class.
-        randomize : bool, default False
+        randomize : bool, default True
             If True, shuffles the data prior to selection - uses a set seed for reproducibility.
         slice_back : bool, default False
             If True and size has a value greater than 0, then grabs selection starting at the last image.
@@ -251,7 +251,7 @@ class MNIST(Dataset[tuple[NDArray[np.float64], int]]):
         corruption: CorruptionStringMap | None = None,
         classes: TClassMap | None = None,
         balance: bool = True,
-        randomize: bool = False,
+        randomize: bool = True,
         slice_back: bool = False,
         verbose: bool = True,
     ) -> None:

dataeval/utils/torch/gmm.py

@@ -0,0 +1,98 @@
+"""
+Adapted for Pytorch from:
+
+Source code derived from Alibi-Detect 0.11.4
+https://github.com/SeldonIO/alibi-detect/tree/v0.11.4
+
+Original code Copyright (c) 2023 Seldon Technologies Ltd
+Licensed under Apache Software License (Apache 2.0)
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+
+from dataeval.utils.gmm import GaussianMixtureModelParams
+
+
+def gmm_params(z: torch.Tensor, gamma: torch.Tensor) -> GaussianMixtureModelParams[torch.Tensor]:
+    """
+    Compute parameters of Gaussian Mixture Model.
+
+    Parameters
+    ----------
+    z : torch.Tensor
+        Observations.
+    gamma : torch.Tensor
+        Mixture probabilities to derive mixture distribution weights from.
+
+    Returns
+    -------
+    GaussianMixtureModelParams(phi, mu, cov, L, log_det_cov)
+        The parameters used to calculate energy.
+    """
+
+    # compute gmm parameters phi, mu and cov
+    N = gamma.shape[0]  # nb of samples in batch
+    sum_gamma = torch.sum(gamma, 0)  # K
+    phi = sum_gamma / N  # K
+    # K x D (D = latent_dim)
+    mu = torch.sum(torch.unsqueeze(gamma, -1) * torch.unsqueeze(z, 1), 0) / torch.unsqueeze(sum_gamma, -1)
+    z_mu = torch.unsqueeze(z, 1) - torch.unsqueeze(mu, 0)  # N x K x D
+    z_mu_outer = torch.unsqueeze(z_mu, -1) * torch.unsqueeze(z_mu, -2)  # N x K x D x D
+
+    # K x D x D
+    cov = torch.sum(torch.unsqueeze(torch.unsqueeze(gamma, -1), -1) * z_mu_outer, 0) / torch.unsqueeze(
+        torch.unsqueeze(sum_gamma, -1), -1
+    )
+
+    # cholesky decomposition of covariance and determinant derivation
+    D = cov.shape[1]
+    eps = 1e-6
+    L = torch.linalg.cholesky(cov + torch.eye(D) * eps)  # K x D x D
+    log_det_cov = 2.0 * torch.sum(torch.log(torch.diagonal(L, dim1=-2, dim2=-1)), 1)  # K
+
+    return GaussianMixtureModelParams(phi, mu, cov, L, log_det_cov)
+
+
+def gmm_energy(
+    z: torch.Tensor,
+    params: GaussianMixtureModelParams[torch.Tensor],
+    return_mean: bool = True,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Compute sample energy from Gaussian Mixture Model.
+
+    Parameters
+    ----------
+    params : GaussianMixtureModelParams
+        The gaussian mixture model parameters.
+    return_mean : bool, default True
+        Take mean across all sample energies in a batch.
+
+    Returns
+    -------
+    sample_energy
+        The sample energy of the GMM.
+    cov_diag
+        The inverse sum of the diagonal components of the covariance matrix.
+    """
+    D = params.cov.shape[1]
+    z_mu = torch.unsqueeze(z, 1) - torch.unsqueeze(params.mu, 0)  # N x K x D
+    z_mu_T = torch.permute(z_mu, dims=[1, 2, 0])  # K x D x N
+    v = torch.linalg.solve_triangular(params.L, z_mu_T, upper=False)  # K x D x D
+
+    # rewrite sample energy in logsumexp format for numerical stability
+    logits = torch.log(torch.unsqueeze(params.phi, -1)) - 0.5 * (
+        torch.sum(torch.square(v), 1) + float(D) * np.log(2.0 * np.pi) + torch.unsqueeze(params.log_det_cov, -1)
+    )  # K x N
+    sample_energy = -torch.logsumexp(logits, 0)  # N
+
+    if return_mean:
+        sample_energy = torch.mean(sample_energy)
+
+    # inverse sum of variances
+    cov_diag = torch.sum(torch.divide(torch.tensor(1), torch.diagonal(params.cov, dim1=-2, dim2=-1)))
+
+    return sample_energy, cov_diag

dataeval/utils/torch/models.py

@@ -2,8 +2,10 @@ from __future__ import annotations
 
 __all__ = ["AriaAutoencoder", "Encoder", "Decoder"]
 
+import math
 from typing import Any
 
+import torch
 import torch.nn as nn
 
 
@@ -136,3 +138,193 @@ class Decoder(nn.Module):
             The reconstructed output tensor.
         """
         return self.decoder(x)
+
+
+class AE(nn.Module):
+    """
+    An autoencoder model with a separate encoder and decoder. Meant to replace the TensorFlow model called AE, which we
+      used as the core of an autoencoder-based OOD detector, i.e. as an argument to OOD_AE().
+
+    Parameters
+    ----------
+    input_shape : tuple[int, int, int]
+        Number of input channels, number of rows, number of columns.() Number of examples per batch will be inferred
+          at runtime.)
+    """
+
+    def __init__(self, input_shape: tuple[int, int, int]) -> None:
+        super().__init__()
+
+        input_dim = math.prod(input_shape)
+
+        # following is lifted from src/dataeval/utils/tensorflow/_internal/utils.py. It makes an odd staircase that is
+        #  basically proportional to the number of numbers in the image to the 0.8 power. '
+        encoding_dim = int(math.pow(2, int(input_dim.bit_length() * 0.8)))
+
+        self.encoder: Encoder_AE = Encoder_AE(input_shape, encoding_dim)
+
+        self.decoder: Decoder_AE = Decoder_AE(input_shape, encoding_dim, self.encoder.post_op_shape)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Perform a forward pass through the encoder and decoder.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor
+
+        Returns
+        -------
+        torch.Tensor
+            The reconstructed output tensor.
+        """
+        x = self.encoder(x)
+        x = self.decoder(x)
+        return x
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Encode the input tensor using the encoder.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor
+
+        Returns
+        -------
+        torch.Tensor
+            The encoded representation of the input tensor.
+        """
+        return self.encoder(x)
+
+
+class Encoder_AE(nn.Module):
+    """
+    A simple encoder to be used in an autoencoder model.
+
+    This is the encoder used to replicate AE, which was a TF function. It consists of a CNN followed by a fully
+      connected layer.
+
+    Parameters
+    ----------
+    channels : int
+        Number of input channels
+
+    input_shape : tuple[int, int, int]
+        number of channels, number of rows, number of columns in input images.
+
+    encoding_dim : the size of the 1D array that emerges from the fully connected layer.
+
+    """
+
+    def __init__(
+        self,
+        input_shape: tuple[int, int, int],
+        encoding_dim: int,
+    ) -> None:
+        super().__init__()
+
+        channels = input_shape[0]
+        nc_in, nc_mid, nc_done = 256, 128, 64
+
+        conv_in = nn.Conv2d(channels, nc_in, 2, stride=1, padding=1)
+        conv_mid = nn.Conv2d(nc_in, nc_mid, 2, stride=1, padding=1)
+        conv_done = nn.Conv2d(nc_mid, nc_done, 2, stride=1)
+
+        self.encoding_ops: nn.Sequential = nn.Sequential(
+            conv_in,
+            nn.LeakyReLU(),
+            nn.MaxPool2d(2),
+            conv_mid,
+            nn.LeakyReLU(),
+            nn.MaxPool2d(2),
+            conv_done,
+        )
+
+        ny, nx = input_shape[1:]
+        self.post_op_shape: tuple[int, int, int] = (nc_done, ny // 4 - 1, nx // 4 - 1)
+        self.flatcon: int = math.prod(self.post_op_shape)
+        self.flatten: nn.Sequential = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(
+                self.flatcon,
+                encoding_dim,
+            ),
+        )
+
+    def forward(self, x: Any) -> Any:
+        """
+        Perform a forward pass through the AE_torch encoder.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor
+
+        Returns
+        -------
+        torch.Tensor
+            The encoded representation of the input tensor.
+        """
+        x = self.encoding_ops(x)
+
+        x = self.flatten(x)
+
+        return x
+
+
+class Decoder_AE(nn.Module):
+    """
+    A simple decoder to be used in an autoencoder model.
+
+    This is the decoder used by the AriaAutoencoder model.
+
+    Parameters
+    ----------
+    channels : int
+        Number of output channels
+    """
+
+    def __init__(
+        self,
+        input_shape: tuple[int, int, int],
+        encoding_dim: int,
+        post_op_shape: tuple[int, int, int],
+    ) -> None:
+        super().__init__()
+
+        self.post_op_shape = post_op_shape
+        self.input_shape = input_shape  # need to store this for use in forward().
+        channels = input_shape[0]
+
+        self.input: nn.Linear = nn.Linear(encoding_dim, math.prod(post_op_shape))
+
+        self.decoder: nn.Sequential = nn.Sequential(
+            nn.ConvTranspose2d(64, 128, 2, stride=1),
+            nn.LeakyReLU(),
+            nn.ConvTranspose2d(128, 256, 2, stride=2),
+            nn.LeakyReLU(),
+            nn.ConvTranspose2d(256, channels, 2, stride=2),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Perform a forward pass through the decoder.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            The encoded tensor.
+
+        Returns
+        -------
+        torch.Tensor
+            The reconstructed output tensor.
+        """
+        x = self.input(x)
+        x = x.reshape((-1, *self.post_op_shape))
+        x = self.decoder(x)
+        x = x.reshape((-1, *self.input_shape))
+        return x

dataeval/utils/torch/trainer.py

@@ -1,15 +1,15 @@
 from __future__ import annotations
 
-__all__ = ["AETrainer"]
-
-from typing import Any
+from typing import Any, Callable
 
 import torch
 import torch.nn as nn
+from numpy.typing import NDArray
 from torch.optim import Adam
-from torch.utils.data import DataLoader, Dataset
+from torch.utils.data import DataLoader, Dataset, TensorDataset
+from tqdm import tqdm
 
-torch.manual_seed(0)
+__all__ = ["AETrainer", "trainer"]
 
 
 def get_images_from_batch(batch: Any) -> Any:
@@ -176,3 +176,82 @@ class AETrainer:
             encodings = torch.vstack((encodings, embeddings)) if len(encodings) else embeddings
 
         return encodings
+
+
+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()

dataeval/utils/torch/utils.py

@@ -3,8 +3,12 @@ from __future__ import annotations
 __all__ = ["read_dataset"]
 
 from collections import defaultdict
-from typing import Any
+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
 
 
@@ -61,3 +65,105 @@ def read_dataset(dataset: Dataset[Any]) -> list[list[Any]]:
             ddict[i].append(d)
 
     return list(ddict.values())
+
+
+def get_device(device: str | torch.device | None = None) -> torch.device:
+    """
+    Instantiates a PyTorch device object.
+
+    Parameters
+    ----------
+    device : str | torch.device | None, default None
+        Either ``None``, a str ('gpu' or 'cpu') indicating the device to choose, or an
+        already instantiated device object. If ``None``, the GPU is selected if it is
+        detected, otherwise the CPU is used as a fallback.
+
+    Returns
+    -------
+    The instantiated device object.
+    """
+    if isinstance(device, torch.device):  # Already a torch device
+        return device
+    else:  # Instantiate device
+        if device is None or device.lower() in ["gpu", "cuda"]:
+            torch_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        else:
+            torch_device = torch.device("cpu")
+    return torch_device
+
+
+def predict_batch(
+    x: NDArray[Any] | torch.Tensor,
+    model: Callable | torch.nn.Module | torch.nn.Sequential,
+    device: torch.device | None = None,
+    batch_size: int = int(1e10),
+    preprocess_fn: Callable | None = None,
+    dtype: type[np.generic] | torch.dtype = np.float32,
+) -> NDArray[Any] | torch.Tensor | tuple[Any, ...]:
+    """
+    Make batch predictions on a model.
+
+    Parameters
+    ----------
+    x : np.ndarray | torch.Tensor
+        Batch of instances.
+    model : Callable | nn.Module | nn.Sequential
+        PyTorch model.
+    device : torch.device | None, default None
+        Device type used. The default None tries to use the GPU and falls back on CPU.
+        Can be specified by passing either torch.device('cuda') or torch.device('cpu').
+    batch_size : int, default 1e10
+        Batch size used during prediction.
+    preprocess_fn : Callable | None, default None
+        Optional preprocessing function for each batch.
+    dtype : np.dtype | torch.dtype, default np.float32
+        Model output type, either a :term:`NumPy` or torch dtype, e.g. np.float32 or torch.float32.
+
+    Returns
+    -------
+    NDArray | torch.Tensor | tuple
+        Numpy array, torch tensor or tuples of those with model outputs.
+    """
+    device = get_device(device)
+    if isinstance(x, np.ndarray):
+        x = torch.from_numpy(x).to(device)
+    n = len(x)
+    n_minibatch = int(np.ceil(n / batch_size))
+    return_np = not isinstance(dtype, torch.dtype)
+    preds = []
+    with torch.no_grad():
+        for i in range(n_minibatch):
+            istart, istop = i * batch_size, min((i + 1) * batch_size, n)
+            x_batch = x[istart:istop]
+            if isinstance(preprocess_fn, Callable):
+                x_batch = preprocess_fn(x_batch)
+
+            preds_tmp = model(x_batch.to(torch.float32).to(device))
+            if isinstance(preds_tmp, (list, tuple)):
+                if len(preds) == 0:  # init tuple with lists to store predictions
+                    preds = tuple([] for _ in range(len(preds_tmp)))
+                for j, p in enumerate(preds_tmp):
+                    if isinstance(p, torch.Tensor):
+                        p = p.cpu()
+                    preds[j].append(p if not return_np or isinstance(p, np.ndarray) else p.numpy())
+            elif isinstance(preds_tmp, (np.ndarray, torch.Tensor)):
+                if isinstance(preds_tmp, torch.Tensor):
+                    preds_tmp = preds_tmp.cpu()
+                if isinstance(preds, tuple):
+                    preds = list(preds)
+                preds.append(
+                    preds_tmp
+                    if not return_np or isinstance(preds_tmp, np.ndarray)  # type: ignore
+                    else preds_tmp.numpy()
+                )
+            else:
+                raise TypeError(
+                    f"Model output type {type(preds_tmp)} not supported. The model \
+                    output type needs to be one of list, tuple, NDArray or \
+                    torch.Tensor."
+                )
+    concat = partial(np.concatenate, axis=0) if return_np else partial(torch.cat, dim=0)
+    out: tuple | np.ndarray | torch.Tensor = (
+        tuple(concat(p) for p in preds) if isinstance(preds, tuple) else concat(preds)  # type: ignore
+    )
+    return out

dataeval/workflows/__init__.py

@@ -4,7 +4,7 @@ Workflows perform a sequence of actions to analyze the dataset and make predicti
 
 from dataeval import _IS_TORCH_AVAILABLE
 
-if _IS_TORCH_AVAILABLE:  # pragma: no cover
+if _IS_TORCH_AVAILABLE:
     from dataeval.workflows.sufficiency import Sufficiency, SufficiencyOutput
 
     __all__ = ["Sufficiency", "SufficiencyOutput"]

{dataeval-0.73.0.dist-info → dataeval-0.74.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dataeval
-Version: 0.73.0
+Version: 0.74.0
 Summary: DataEval provides a simple interface to characterize image data and its impact on model performance across classification and object-detection tasks
 Home-page: https://dataeval.ai/
 License: MIT
@@ -23,7 +23,6 @@ Classifier: Topic :: Scientific/Engineering
 Provides-Extra: all
 Provides-Extra: tensorflow
 Provides-Extra: torch
-Requires-Dist: hdbscan (>=0.8.36)
 Requires-Dist: markupsafe (<3.0.2) ; extra == "tensorflow" or extra == "all"
 Requires-Dist: matplotlib ; extra == "torch" or extra == "all"
 Requires-Dist: numpy (>1.24.3)