oodeel 0.4.0__py3-none-any.whl

oodeel/utils/tf_training_tools.py

@@ -0,0 +1,219 @@
+# -*- coding: utf-8 -*-
+# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.layers import Conv2D
+from tensorflow.keras.layers import Dense
+from tensorflow.keras.layers import Dropout
+from tensorflow.keras.layers import Flatten
+from tensorflow.keras.layers import MaxPooling2D
+from tensorflow.keras.models import Sequential
+
+from ..datasets.tf_data_handler import TFDataHandler
+from ..types import List
+from ..types import Optional
+from ..types import Union
+
+
+def get_toy_mlp(input_shape: tuple, num_classes: int) -> tf.keras.Model:
+    """Basic keras MLP classifier for toy datasets.
+
+    Args:
+        input_shape (tuple): Input data shape.
+        num_classes (int): Number of classes for the classification task.
+
+    Returns:
+        tf.keras.Model: model
+    """
+    return tf.keras.models.Sequential(
+        [
+            tf.keras.layers.Input(shape=input_shape),
+            tf.keras.layers.Flatten(),
+            tf.keras.layers.Dense(300, activation="relu"),
+            tf.keras.layers.Dense(150, activation="relu"),
+            tf.keras.layers.Dense(num_classes, activation="softmax"),
+        ]
+    )
+
+
+def get_toy_keras_convnet(num_classes: int) -> tf.keras.Model:
+    """Basic keras convolutional classifier for toy datasets.
+
+    Args:
+        num_classes (int): Number of classes for the classification task.
+
+    Returns:
+        tf.keras.Model: model
+    """
+    return Sequential(
+        [
+            Conv2D(32, kernel_size=(3, 3), activation="relu"),
+            MaxPooling2D(pool_size=(2, 2)),
+            Conv2D(64, kernel_size=(3, 3), activation="relu"),
+            MaxPooling2D(pool_size=(2, 2)),
+            Flatten(),
+            Dropout(0.5),
+            Dense(num_classes, activation="softmax"),
+        ]
+    )
+
+
+def train_tf_model(
+    train_data: tf.data.Dataset,
+    model: Union[tf.keras.Model, str],
+    input_shape: tuple = None,
+    num_classes: int = None,
+    batch_size: int = 128,
+    epochs: int = 50,
+    loss: str = "sparse_categorical_crossentropy",
+    optimizer: str = "adam",
+    lr_scheduler: Optional[str] = None,
+    learning_rate: float = 1e-3,
+    metrics: List[str] = ["accuracy"],
+    imagenet_pretrained: bool = False,
+    validation_data: Optional[tf.data.Dataset] = None,
+    save_dir: Optional[str] = None,
+    save_best_only: bool = True,
+) -> tf.keras.Model:
+    """Loads a model from tensorflow.python.keras.applications.
+    If the dataset is different from imagenet, trains on provided dataset.
+
+    Args:
+        train_data (tf.data.Dataset): training dataset.
+        model (Union[tf.keras.Model, str]): if a string is provided, must be a model
+            from tf.keras.applications or "toy_convnet" or "toy_mlp"
+        input_shape (tuple): Shape of the input images.
+        num_classes (int): Number of output classes.
+        batch_size (int, optional): Defaults to 128.
+        epochs (int, optional): Defaults to 50.
+        loss (str, optional): Defaults to "sparse_categorical_crossentropy".
+        optimizer (str, optional): Defaults to "adam".
+        lr_scheduler (str, optional): ("cosine" | "steps" | None). Defaults to None.
+        learning_rate (float, optional): Defaults to 1e-3.
+        metrics (List[str], optional): Validation metrics. Defaults to ["accuracy"].
+        imagenet_pretrained (bool, optional): Load a model pretrained on imagenet or
+            not. Defaults to False.
+        validation_data (Optional[tf.data.Dataset], optional): Defaults to None.
+        save_dir (Optional[str], optional): Directory to save the model.
+            Defaults to None.
+        save_best_only (bool): If False, saved model will be the last one. Defaults to
+            True.
+
+    Returns:
+        tf.keras.Model: Trained model
+    """
+    # get data infos from dataset
+    if isinstance(train_data.element_spec, dict):
+        input_id = "image"
+        label_id = "label"
+    else:
+        input_id = 0
+        label_id = -1
+    if input_shape is None:
+        input_shape = TFDataHandler.get_feature_shape(train_data, input_id)
+    if num_classes is None:
+        classes = TFDataHandler.get_feature(train_data, label_id).unique()
+        num_classes = len(list(classes.as_numpy_iterator()))
+
+    # prepare model
+    if isinstance(model, tf.keras.Model):
+        pass
+    elif isinstance(model, str):
+        if model == "toy_convnet":
+            model = get_toy_keras_convnet(num_classes)
+        elif model == "toy_mlp":
+            model = get_toy_mlp(input_shape, num_classes)
+        else:
+            weights = "imagenet" if imagenet_pretrained else None
+            backbone = getattr(tf.keras.applications, model)(
+                include_top=False, weights=weights, input_shape=input_shape
+            )
+
+            features = tf.keras.layers.Flatten()(backbone.layers[-1].output)
+            output = tf.keras.layers.Dense(
+                num_classes,
+                activation="softmax",
+            )(features)
+            model = tf.keras.Model(backbone.layers[0].input, output)
+
+    n_samples = TFDataHandler.get_dataset_length(train_data)
+
+    # Prepare callbacks
+    model_checkpoint_callback = []
+
+    if save_dir is not None:
+        checkpoint_filepath = save_dir
+        model_checkpoint_callback.append(
+            tf.keras.callbacks.ModelCheckpoint(
+                filepath=checkpoint_filepath,
+                save_weights_only=True,
+                monitor="val_accuracy",
+                mode="max",
+                save_best_only=save_best_only,
+            )
+        )
+
+    if len(model_checkpoint_callback) == 0:
+        model_checkpoint_callback = None
+
+    # optimizer
+    decay_steps = int(epochs * n_samples / batch_size)
+    if lr_scheduler == "cosine":
+        learning_rate_fn = tf.keras.experimental.CosineDecay(
+            learning_rate, decay_steps=decay_steps
+        )
+    elif lr_scheduler == "steps":
+        values = list(learning_rate * np.array([1, 0.1, 0.01]))
+        boundaries = list(np.round(decay_steps * np.array([1 / 3, 2 / 3])).astype(int))
+        learning_rate_fn = tf.keras.optimizers.schedules.PiecewiseConstantDecay(
+            boundaries, values
+        )
+    else:
+        learning_rate_fn = learning_rate
+
+    config = {
+        "class_name": optimizer,
+        "config": {
+            "learning_rate": learning_rate_fn,
+        },
+    }
+
+    if optimizer == "SGD":
+        config["config"]["momentum"] = 0.9
+        config["config"]["decay"] = 5e-4
+
+    keras_optimizer = tf.keras.optimizers.get(config)
+
+    model.compile(loss=loss, optimizer=keras_optimizer, metrics=metrics)
+
+    model.fit(
+        train_data,
+        validation_data=validation_data,
+        epochs=epochs,
+        callbacks=model_checkpoint_callback,
+    )
+
+    if save_dir is not None:
+        model.load_weights(save_dir)
+        model.save(save_dir)
+    return model

oodeel/utils/torch_operator.py

@@ -0,0 +1,292 @@
+# -*- coding: utf-8 -*-
+# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+from typing import List
+from typing import Optional
+
+import numpy as np
+import torch
+
+from ..types import Callable
+from ..types import TensorType
+from ..types import Union
+from .general_utils import is_from
+from .operator import Operator
+
+
+def sanitize_input(tensor_arg_func: Callable):
+    """ensures the decorated function receives a torch.Tensor"""
+
+    def wrapper(obj, tensor, *args, **kwargs):
+        if isinstance(tensor, torch.Tensor):
+            pass
+        elif is_from(tensor, "tensorflow"):
+            tensor = torch.Tensor(tensor.numpy())
+        else:
+            tensor = torch.Tensor(tensor)
+
+        return tensor_arg_func(obj, tensor, *args, **kwargs)
+
+    return wrapper
+
+
+class TorchOperator(Operator):
+    """Class to handle torch operations with a unified API"""
+
+    def __init__(self, model: Optional[torch.nn.Module] = None):
+        if model is not None:
+            self._device = next(model.parameters()).device
+        else:
+            self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    @staticmethod
+    def softmax(tensor: TensorType) -> torch.Tensor:
+        """Softmax function along the last dimension"""
+        return torch.nn.functional.softmax(tensor, dim=-1)
+
+    @staticmethod
+    def argmax(tensor: TensorType, dim: Optional[int] = None) -> torch.Tensor:
+        """Argmax function"""
+        return torch.argmax(tensor, dim=dim)
+
+    @staticmethod
+    def max(
+        tensor: TensorType, dim: Optional[int] = None, keepdim: Optional[bool] = False
+    ) -> torch.Tensor:
+        """Max function"""
+        if dim is None:
+            return torch.max(tensor)
+        else:
+            return torch.max(tensor, dim, keepdim=keepdim)[0]
+
+    @staticmethod
+    def min(
+        tensor: TensorType, dim: Optional[int] = None, keepdim: bool = False
+    ) -> torch.Tensor:
+        """Min function"""
+        if dim is None:
+            return torch.min(tensor)
+        else:
+            return torch.min(tensor, dim, keepdim=keepdim)[0]
+
+    @staticmethod
+    def one_hot(tensor: TensorType, num_classes: int) -> torch.Tensor:
+        """One hot function"""
+        return torch.nn.functional.one_hot(tensor, num_classes)
+
+    @staticmethod
+    def sign(tensor: TensorType) -> torch.Tensor:
+        """Sign function"""
+        return torch.sign(tensor)
+
+    @staticmethod
+    def CrossEntropyLoss(reduction: str = "mean"):
+        """Cross Entropy Loss from logits"""
+
+        def sanitized_ce_loss(inputs, targets):
+            return torch.nn.CrossEntropyLoss(reduction=reduction)(inputs, targets)
+
+        return sanitized_ce_loss
+
+    @staticmethod
+    def norm(tensor: TensorType, dim: Optional[int] = None) -> torch.Tensor:
+        """Tensor Norm"""
+        return torch.norm(tensor, dim=dim)
+
+    @staticmethod
+    def matmul(tensor_1: TensorType, tensor_2: TensorType) -> torch.Tensor:
+        """Matmul operation"""
+        return torch.matmul(tensor_1, tensor_2)
+
+    @staticmethod
+    def convert_from_tensorflow(tensor: TensorType) -> torch.Tensor:
+        """Convert a tensorflow tensor into a torch tensor
+
+        Used when using a pytorch model on a dataset loaded from tensorflow datasets
+        """
+        return torch.Tensor(tensor.numpy())
+
+    @staticmethod
+    def convert_to_numpy(tensor: TensorType) -> np.ndarray:
+        """Convert tensor into a np.ndarray"""
+        if not isinstance(tensor, np.ndarray):
+            if tensor.device != "cpu":
+                tensor = tensor.to("cpu")
+            return tensor.detach().numpy()
+        return tensor
+
+    @staticmethod
+    def gradient(func: Callable, inputs: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+        """Compute gradients for a batch of samples.
+
+        Args:
+            func (Callable): Function used for computing gradient. Must be built with
+                torch differentiable operations only, and return a scalar.
+            inputs (torch.Tensor): Input tensor wrt which the gradients are computed
+            *args: Additional Args for func.
+            **kwargs: Additional Kwargs for func.
+
+        Returns:
+            torch.Tensor: Gradients computed, with the same shape as the inputs.
+        """
+        inputs.requires_grad_(True)
+        outputs = func(inputs, *args, **kwargs)
+        gradients = torch.autograd.grad(outputs, inputs)
+        inputs.requires_grad_(False)
+        return gradients[0]
+
+    @staticmethod
+    def stack(tensors: List[TensorType], dim: int = 0) -> torch.Tensor:
+        "Stack tensors along a new dimension"
+        return torch.stack(tensors, dim)
+
+    @staticmethod
+    def cat(tensors: List[TensorType], dim: int = 0) -> torch.Tensor:
+        "Concatenate tensors in a given dimension"
+        return torch.cat(tensors, dim)
+
+    @staticmethod
+    def mean(tensor: TensorType, dim: Optional[int] = None) -> torch.Tensor:
+        "Mean function"
+        if dim is None:
+            return torch.mean(tensor)
+        else:
+            return torch.mean(tensor, dim)
+
+    @staticmethod
+    def flatten(tensor: TensorType) -> torch.Tensor:
+        "Flatten function"
+        # Flatten the features to 2D (n_batch, n_features)
+        return tensor.view(tensor.size(0), -1)
+
+    def from_numpy(self, arr: np.ndarray) -> torch.Tensor:
+        "Convert a NumPy array to a tensor"
+        # TODO change dtype
+        return torch.tensor(arr).to(self._device)
+
+    @staticmethod
+    def t(tensor: TensorType) -> torch.Tensor:
+        "Transpose function for tensor of rank 2"
+        return tensor.t()
+
+    @staticmethod
+    def permute(tensor: TensorType, dims) -> torch.Tensor:
+        "Transpose function for tensor of rank 2"
+        return torch.permute(tensor, dims)
+
+    @staticmethod
+    def diag(tensor: TensorType) -> torch.Tensor:
+        "Diagonal function: return the diagonal of a 2D tensor"
+        return tensor.diag()
+
+    @staticmethod
+    def reshape(tensor: TensorType, shape: List[int]) -> torch.Tensor:
+        "Reshape function"
+        return tensor.view(*shape)
+
+    @staticmethod
+    def equal(tensor: TensorType, other: Union[TensorType, int, float]) -> torch.Tensor:
+        "Computes element-wise equality"
+        return torch.eq(tensor, other)
+
+    @staticmethod
+    def pinv(tensor: TensorType) -> torch.Tensor:
+        "Computes the pseudoinverse (Moore-Penrose inverse) of a matrix."
+        return torch.linalg.pinv(tensor)
+
+    @staticmethod
+    def eigh(tensor: TensorType) -> torch.Tensor:
+        "Computes the eigen decomposition of a self-adjoint matrix."
+        eigval, eigvec = torch.linalg.eigh(tensor)
+        return eigval, eigvec
+
+    @staticmethod
+    def quantile(tensor: TensorType, q: float, dim: int = None) -> torch.Tensor:
+        "Computes the quantile of a tensor's components. q in (0,1)"
+        if dim is None:
+            # keep the 16 millions first elements (see torch.quantile issue:
+            # https://github.com/pytorch/pytorch/issues/64947)
+            tensor_flatten = tensor.view(-1)[:16_000_000]
+            return torch.quantile(tensor_flatten, q).item()
+        else:
+            return torch.quantile(tensor, q, dim)
+
+    @staticmethod
+    def relu(tensor: TensorType) -> torch.Tensor:
+        "Apply relu to a tensor"
+        return torch.nn.functional.relu(tensor)
+
+    @staticmethod
+    def einsum(equation: str, *tensors: TensorType) -> torch.Tensor:
+        "Computes the einsum between tensors following equation"
+        return torch.einsum(equation, *tensors)
+
+    @staticmethod
+    def tril(tensor: TensorType, diagonal: int = 0) -> torch.Tensor:
+        "Set the upper triangle of the matrix formed by the last two dimensions of"
+        "tensor to zero"
+        return torch.tril(tensor, diagonal)
+
+    @staticmethod
+    def sum(tensor: TensorType, dim: Union[tuple, list, int] = None) -> torch.Tensor:
+        "sum along dim"
+        if dim is None:
+            return torch.sum(tensor)
+        return torch.sum(tensor, dim)
+
+    @staticmethod
+    def unsqueeze(tensor: TensorType, dim: int) -> torch.Tensor:
+        "unsqueeze along dim"
+        return torch.unsqueeze(tensor, dim)
+
+    @staticmethod
+    def squeeze(tensor: TensorType, dim: int = None) -> torch.Tensor:
+        "squeeze along dim"
+
+        if dim is None:
+            return torch.squeeze(tensor)
+
+        return torch.squeeze(tensor, dim)
+
+    @staticmethod
+    def abs(tensor: TensorType) -> torch.Tensor:
+        "compute absolute value"
+        return torch.abs(tensor)
+
+    @staticmethod
+    def where(
+        condition: TensorType,
+        input: Union[TensorType, float],
+        other: Union[TensorType, float],
+    ) -> torch.Tensor:
+        "Applies where function , to condition"
+        return torch.where(condition, input, other)
+
+    @staticmethod
+    def avg_pool_2d(tensor: TensorType) -> torch.Tensor:
+        """Perform avg pool in 2d as in torch.nn.functional.adaptive_avg_pool2d"""
+        return torch.mean(tensor, dim=(-2, -1))
+
+    @staticmethod
+    def log(tensor: TensorType) -> torch.Tensor:
+        """Perform log"""
+        return torch.log(tensor)