Merge pull request #1360 from Trusted-AI/dev_jax_estimator

Implement a classifier on Jax framework

Author: beat-buesser

.github/workflows/ci-lingvo.yml

@@ -47,7 +47,7 @@ jobs:
           sudo apt-get update
           sudo apt-get -y -q install ffmpeg libavcodec-extra
           python -m pip install --upgrade pip setuptools wheel
-          pip install -q -r <(sed '/^scipy/d;/^matplotlib/d;/^pandas/d;/^statsmodels/d;/^numba/d' requirements_test.txt)
+          pip install -q -r <(sed '/^scipy/d;/^matplotlib/d;/^pandas/d;/^statsmodels/d;/^numba/d;/^jax/d' requirements_test.txt)
           pip install scipy==1.5.4
           pip install matplotlib==3.3.4
           pip install pandas==1.1.5
@@ -58,6 +58,7 @@ jobs:
           pip install lingvo==${{ matrix.lingvo }}
           pip install tensorflow-addons==0.9.1
           pip install model-pruning-google-research==0.0.3
+          pip install jax[cpu]==0.2.17
           pip list
       - name: Run ${{ matrix.name }} Tests
         run: ./run_tests.sh ${{ matrix.framework }}

art/experimental/__init__.py

@@ -0,0 +1,4 @@
+"""
+This module contains the experimental Estimator API.
+"""
+from art.experimental.estimators.jax import JaxEstimator

art/experimental/estimators/__init__.py

@@ -0,0 +1,4 @@
+"""
+Experimental classifiers.
+"""
+from art.experimental.estimators.classification.jax import JaxClassifier

art/experimental/estimators/classification/__init__.py

@@ -0,0 +1,302 @@
+# MIT License
+#
+# Copyright (C) The Adversarial Robustness Toolbox (ART) Authors 2021
+#
+# 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.
+"""
+This module implements the classifier `JaxClassifier` for Jax models.
+"""
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import logging
+import random
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union, TYPE_CHECKING
+
+import numpy as np
+
+from art.estimators.classification.classifier import (
+    ClassGradientsMixin,
+    ClassifierMixin,
+)
+from art.experimental.estimators.jax import JaxEstimator
+
+if TYPE_CHECKING:
+    from art.utils import CLIP_VALUES_TYPE, PREPROCESSING_TYPE
+    from art.data_generators import DataGenerator
+    from art.defences.preprocessor import Preprocessor
+    from art.defences.postprocessor import Postprocessor
+
+logger = logging.getLogger(__name__)
+
+
+class JaxClassifier(ClassGradientsMixin, ClassifierMixin, JaxEstimator):  # lgtm [py/missing-call-to-init]
+    """
+    This class implements a classifier with the Jax framework.
+    """
+
+    estimator_params = (
+        JaxEstimator.estimator_params
+        + ClassifierMixin.estimator_params
+        + [
+            "predict_func",
+            "loss_func",
+            "update_func",
+        ]
+    )
+
+    def __init__(
+        self,
+        model: List,
+        predict_func: Callable,
+        loss_func: Callable,
+        update_func: Callable,
+        input_shape: Tuple[int, ...],
+        nb_classes: int,
+        channels_first: bool = False,
+        clip_values: Optional["CLIP_VALUES_TYPE"] = None,
+        preprocessing_defences: Union["Preprocessor", List["Preprocessor"], None] = None,
+        postprocessing_defences: Union["Postprocessor", List["Postprocessor"], None] = None,
+        preprocessing: "PREPROCESSING_TYPE" = (0.0, 1.0),
+    ) -> None:
+        """
+        Initialization specifically for the Jax-based implementation.
+
+        :param model: Jax model, represented as a list of model parameters.
+        :param predict_func: A function used to predict model output given the model and the input.
+        :param loss_func: The loss function for which to compute gradients for training.
+        :param update_func: The update function for which to train the model.
+        :param input_shape: The shape of one input instance.
+        :param nb_classes: The number of classes of the model.
+        :param channels_first: Set channels first or last.
+        :param clip_values: Tuple of the form `(min, max)` of floats or `np.ndarray` representing the minimum and
+               maximum values allowed for features. If floats are provided, these will be used as the range of all
+               features. If arrays are provided, each value will be considered the bound for a feature, thus
+               the shape of clip values needs to match the total number of features.
+        :param preprocessing_defences: Preprocessing defence(s) to be applied by the classifier.
+        :param postprocessing_defences: Postprocessing defence(s) to be applied by the classifier.
+        :param preprocessing: Tuple of the form `(subtrahend, divisor)` of floats or `np.ndarray` of values to be
+               used for data preprocessing. The first value will be subtracted from the input. The input will then
+               be divided by the second one.
+        """
+        super().__init__(
+            model=model,
+            clip_values=clip_values,
+            channels_first=channels_first,
+            preprocessing_defences=preprocessing_defences,
+            postprocessing_defences=postprocessing_defences,
+            preprocessing=preprocessing,
+        )
+
+        self._predict_func = predict_func
+        self._loss_func = loss_func
+        self._update_func = update_func
+        self._nb_classes = nb_classes
+        self._input_shape = input_shape
+
+    @property
+    def model(self) -> List:
+        return self._model
+
+    @property
+    def input_shape(self) -> Tuple[int, ...]:
+        """
+        Return the shape of one input sample.
+
+        :return: Shape of one input sample.
+        """
+        return self._input_shape  # type: ignore
+
+    @property
+    def predict_func(self) -> Callable:
+        """
+        Return the predict function.
+
+        :return: The predict function.
+        """
+        return self._predict_func
+
+    @property
+    def loss_func(self) -> Callable:
+        """
+        Return the loss function.
+
+        :return: The loss function.
+        """
+        return self._loss_func
+
+    @property
+    def update_func(self) -> Callable:
+        """
+        Return the update function.
+
+        :return: The update function.
+        """
+        return self._update_func
+
+    def predict(self, x: np.ndarray, batch_size: int = 128, **kwargs) -> np.ndarray:
+        """
+        Perform prediction for a batch of inputs.
+
+        :param x: Input samples.
+        :param batch_size: Size of batches.
+        :return: Array of predictions of shape `(nb_inputs, nb_classes)`.
+        """
+        # Apply preprocessing
+        x_preprocessed, _ = self._apply_preprocessing(x, y=None, fit=False)
+
+        results_list = []
+
+        # Run prediction with batch processing
+        num_batch = int(np.ceil(len(x_preprocessed) / float(batch_size)))
+        for m in range(num_batch):
+            # Batch indexes
+            begin, end = (
+                m * batch_size,
+                min((m + 1) * batch_size, x_preprocessed.shape[0]),
+            )
+
+            output = self.predict_func(self.model, x_preprocessed[begin:end])
+            results_list.append(output)
+
+        results = np.vstack(results_list)
+
+        # Apply postprocessing
+        predictions = self._apply_postprocessing(preds=results, fit=False)
+
+        return predictions
+
+    def fit(self, x: np.ndarray, y: np.ndarray, batch_size: int = 128, nb_epochs: int = 10, **kwargs) -> None:
+        """
+        Fit the classifier on the training set `(x, y)`.
+
+        :param x: Training data.
+        :param y: Target values.
+        :param batch_size: Size of batches.
+        :param nb_epochs: Number of epochs to use for training.
+        :param kwargs: Dictionary of framework-specific arguments. This parameter is not currently supported for PyTorch
+               and providing it takes no effect.
+        """
+        # Apply preprocessing
+        x_preprocessed, y_preprocessed = self._apply_preprocessing(x, y, fit=True)
+
+        num_batch = int(np.ceil(len(x_preprocessed) / float(batch_size)))
+        ind = np.arange(len(x_preprocessed))
+
+        # Start training
+        for _ in range(nb_epochs):
+            # Shuffle the examples
+            random.shuffle(ind)
+
+            # Train for one epoch
+            for m in range(num_batch):
+                i_batch = x_preprocessed[ind[m * batch_size : (m + 1) * batch_size]]
+                o_batch = y_preprocessed[ind[m * batch_size : (m + 1) * batch_size]]
+                self._model = self.update_func(self.model, i_batch, o_batch)
+
+    def loss_gradient(self, x: np.ndarray, y: np.ndarray, **kwargs) -> np.ndarray:
+        """
+        Compute the gradient of the loss function w.r.t. `x`.
+
+        :param x: Sample input with shape as expected by the model.
+        :param y: Target values.
+        :return: Array of gradients of the same shape as `x`.
+        """
+        from jax import grad
+
+        # Apply preprocessing
+        x_preprocessed, y_preprocessed = self._apply_preprocessing(x, y, fit=False)
+
+        # Compute gradients
+        grads = grad(self.loss_func, argnums=(0, 1))(self.model, x_preprocessed, y_preprocessed)[1]
+
+        assert grads.shape == x.shape
+
+        return grads.copy()
+
+    def fit_generator(self, generator: "DataGenerator", nb_epochs: int = 20, **kwargs) -> None:
+        """
+        Fit the classifier using the generator that yields batches as specified.
+
+        :param generator: Batch generator providing `(x, y)` for each epoch.
+        :param nb_epochs: Number of epochs to use for training.
+        :param kwargs: Dictionary of framework-specific arguments.
+        """
+        raise NotImplementedError
+
+    def class_gradient(  # pylint: disable=W0221
+        self, x: np.ndarray, label: Union[int, List[int], None] = None, **kwargs
+    ) -> np.ndarray:
+        """
+        Compute per-class derivatives w.r.t. `x`.
+
+        :param x: Sample input with shape as expected by the model.
+        :param label: Index of a specific per-class derivative. If an integer is provided, the gradient of that class
+                      output is computed for all samples. If multiple values as provided, the first dimension should
+                      match the batch size of `x`, and each value will be used as target for its corresponding sample in
+                      `x`. If `None`, then gradients for all classes will be computed for each sample.
+        :return: Array of gradients of input features w.r.t. each class in the form
+                 `(batch_size, nb_classes, input_shape)` when computing for all classes, otherwise shape becomes
+                 `(batch_size, 1, input_shape)` when `label` parameter is specified.
+        """
+        raise NotImplementedError
+
+    def get_activations(
+        self,
+        x: np.ndarray,
+        layer: Optional[Union[int, str]] = None,
+        batch_size: int = 128,
+        framework: bool = False,
+    ) -> np.ndarray:
+        """
+        Return the output of the specified layer for input `x`. `layer` is specified by layer index (between 0 and
+        `nb_layers - 1`) or by name. The number of layers can be determined by counting the results returned by
+        calling `layer_names`.
+
+        :param x: Input for computing the activations.
+        :param layer: Layer for computing the activations
+        :param batch_size: Size of batches.
+        :param framework: If true, return the intermediate tensor representation of the activation.
+        :return: The output of `layer`, where the first dimension is the batch size corresponding to `x`.
+        """
+        raise NotImplementedError
+
+    def save(self, filename: str, path: Optional[str] = None) -> None:
+        """
+        Save a model to file in the format specific to the backend framework.
+
+        :param filename: Name of the file where to store the model.
+        :param path: Path of the folder where to store the model. If no path is specified, the model will be stored in
+                     the default data location of the library `ART_DATA_PATH`.
+        """
+        raise NotImplementedError
+
+    def __getstate__(self) -> Dict[str, Any]:
+        """
+        Use to ensure `JaxClassifier` can be pickled.
+
+        :return: State dictionary with instance parameters.
+        """
+        raise NotImplementedError
+
+    def __setstate__(self, state: Dict[str, Any]) -> None:
+        """
+        Use to ensure `JaxClassifier` can be unpickled.
+
+        :param state: State dictionary with instance parameters to restore.
+        """
+        raise NotImplementedError
+
+    def compute_loss(self, x: np.ndarray, y: np.ndarray, **kwargs) -> np.ndarray:
+        raise NotImplementedError