Merge branch 'dev_1.15.0' into object-detector-resize

Author: beat-buesser

art/attacks/evasion/auto_conjugate_gradient.py

@@ -224,7 +224,8 @@ def __call__(self, y_true: tf.Tensor, y_pred: tf.Tensor, *args, **kwargs) -> tf.
                     nb_classes=estimator.nb_classes,
                     input_shape=estimator.input_shape,
                     loss_object=_loss_object_tf,
-                    train_step=estimator._train_step,
+                    optimizer=estimator.optimizer,
+                    train_step=estimator.train_step,
                     channels_first=estimator.channels_first,
                     clip_values=estimator.clip_values,
                     preprocessing_defences=estimator.preprocessing_defences,

art/attacks/evasion/auto_projected_gradient_descent.py

@@ -203,7 +203,8 @@ def __call__(self, y_true: tf.Tensor, y_pred: tf.Tensor, *args, **kwargs) -> tf.
                     nb_classes=estimator.nb_classes,
                     input_shape=estimator.input_shape,
                     loss_object=_loss_object_tf,
-                    train_step=estimator._train_step,
+                    optimizer=estimator.optimizer,
+                    train_step=estimator.train_step,
                     channels_first=estimator.channels_first,
                     clip_values=estimator.clip_values,
                     preprocessing_defences=estimator.preprocessing_defences,

art/attacks/evasion/brendel_bethge.py

@@ -2055,7 +2055,8 @@ def logits_difference(y_true, y_pred):
                 nb_classes=estimator.nb_classes,
                 input_shape=estimator.input_shape,
                 loss_object=self._loss_object,
-                train_step=estimator._train_step,
+                optimizer=estimator.optimizer,
+                train_step=estimator.train_step,
                 channels_first=estimator.channels_first,
                 clip_values=estimator.clip_values,
                 preprocessing_defences=estimator.preprocessing_defences,

art/estimators/certification/derandomized_smoothing/tensorflow.py

@@ -68,6 +68,7 @@ def __init__(
         logits: bool,
         input_shape: Tuple[int, ...],
         loss_object: Optional["tf.Tensor"] = None,
+        optimizer: Optional["tf.keras.optimizers.Optimizer"] = None,
         train_step: Optional[Callable] = None,
         channels_first: bool = False,
         clip_values: Optional["CLIP_VALUES_TYPE"] = None,
@@ -88,8 +89,12 @@ def __init__(
         :param logits: if the model returns logits or normalized probabilities
         :param input_shape: Shape of one input for the classifier, e.g. for MNIST input_shape=(28, 28, 1).
         :param loss_object: The loss function for which to compute gradients. This parameter is applied for training
-            the model and computing gradients of the loss w.r.t. the input.
-        :param train_step: A function that applies a gradient update to the trainable variables.
+               the model and computing gradients of the loss w.r.t. the input.
+        :param optimizer: The optimizer used to train the classifier.
+        :param train_step: A function that applies a gradient update to the trainable variables with signature
+               `train_step(model, images, labels)`. This will override the default training loop that uses the
+               provided `loss_object` and `optimizer` parameters. It is recommended to use the `@tf.function`
+               decorator, if possible, for efficient training.
         :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
@@ -106,6 +111,7 @@ def __init__(
             nb_classes=nb_classes,
             input_shape=input_shape,
             loss_object=loss_object,
+            optimizer=optimizer,
             train_step=train_step,
             channels_first=channels_first,
             clip_values=clip_values,
@@ -144,10 +150,30 @@ def fit(self, x: np.ndarray, y: np.ndarray, batch_size: int = 128, nb_epochs: in
                        "scheduler" which is an optional function that will be called at the end of every
                        epoch to adjust the learning rate.
         """
+        import tensorflow as tf
+
         if self._train_step is None:  # pragma: no cover
-            raise TypeError(
-                "The training function `train_step` is required for fitting a model but it has not been " "defined."
-            )
+            if self._loss_object is None:  # pragma: no cover
+                raise TypeError(
+                    "A loss function `loss_object` or training function `train_step` is required for fitting the "
+                    "model, but it has not been defined."
+                )
+            if self._optimizer is None:  # pragma: no cover
+                raise ValueError(
+                    "An optimizer `optimizer` or training function `train_step` is required for fitting the "
+                    "model, but it has not been defined."
+                )
+
+            @tf.function
+            def train_step(model, images, labels):
+                with tf.GradientTape() as tape:
+                    predictions = model(images, training=True)
+                    loss = self.loss_object(labels, predictions)
+                gradients = tape.gradient(loss, model.trainable_variables)
+                self.optimizer.apply_gradients(zip(gradients, model.trainable_variables))
+
+        else:
+            train_step = self._train_step
 
         scheduler = kwargs.get("scheduler")
 
@@ -167,7 +193,7 @@ def fit(self, x: np.ndarray, y: np.ndarray, batch_size: int = 128, nb_epochs: in
                 i_batch = np.copy(x_preprocessed[ind[m * batch_size : (m + 1) * batch_size]])
                 labels = y_preprocessed[ind[m * batch_size : (m + 1) * batch_size]]
                 images = self.ablator.forward(i_batch)
-                self._train_step(self.model, images, labels)
+                train_step(self.model, images, labels)
 
             if scheduler is not None:
                 scheduler(epoch)

art/estimators/certification/randomized_smoothing/tensorflow.py

@@ -60,6 +60,7 @@ def __init__(
         nb_classes: int,
         input_shape: Tuple[int, ...],
         loss_object: Optional["tf.Tensor"] = None,
+        optimizer: Optional["tf.keras.optimizers.Optimizer"] = None,
         train_step: Optional[Callable] = None,
         channels_first: bool = False,
         clip_values: Optional["CLIP_VALUES_TYPE"] = None,
@@ -78,8 +79,12 @@ def __init__(
         :param nb_classes: the number of classes in the classification task.
         :param input_shape: Shape of one input for the classifier, e.g. for MNIST input_shape=(28, 28, 1).
         :param loss_object: The loss function for which to compute gradients. This parameter is applied for training
-            the model and computing gradients of the loss w.r.t. the input.
-        :param train_step: A function that applies a gradient update to the trainable variables.
+               the model and computing gradients of the loss w.r.t. the input.
+        :param optimizer: The optimizer used to train the classifier.
+        :param train_step: A function that applies a gradient update to the trainable variables with signature
+               `train_step(model, images, labels)`. This will override the default training loop that uses the
+               provided `loss_object` and `optimizer` parameters. It is recommended to use the `@tf.function`
+               decorator, if possible, for efficient training.
         :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
@@ -105,6 +110,7 @@ def __init__(
             nb_classes=nb_classes,
             input_shape=input_shape,
             loss_object=loss_object,
+            optimizer=optimizer,
             train_step=train_step,
             channels_first=channels_first,
             clip_values=clip_values,
@@ -137,9 +143,27 @@ def fit(self, x: np.ndarray, y: np.ndarray, batch_size: int = 128, nb_epochs: in
         import tensorflow as tf
 
         if self._train_step is None:  # pragma: no cover
-            raise TypeError(
-                "The training function `train_step` is required for fitting a model but it has not been " "defined."
-            )
+            if self._loss_object is None:  # pragma: no cover
+                raise TypeError(
+                    "A loss function `loss_object` or training function `train_step` is required for fitting the "
+                    "model, but it has not been defined."
+                )
+            if self._optimizer is None:  # pragma: no cover
+                raise ValueError(
+                    "An optimizer `optimizer` or training function `train_step` is required for fitting the "
+                    "model, but it has not been defined."
+                )
+
+            @tf.function
+            def train_step(model, images, labels):
+                with tf.GradientTape() as tape:
+                    predictions = model(images, training=True)
+                    loss = self.loss_object(labels, predictions)
+                gradients = tape.gradient(loss, model.trainable_variables)
+                self.optimizer.apply_gradients(zip(gradients, model.trainable_variables))
+
+        else:
+            train_step = self._train_step
 
         y = check_and_transform_label_format(y, nb_classes=self.nb_classes)
 
@@ -156,7 +180,7 @@ def fit(self, x: np.ndarray, y: np.ndarray, batch_size: int = 128, nb_epochs: in
             for images, labels in train_ds:
                 # Add random noise for randomized smoothing
                 images += tf.random.normal(shape=images.shape, mean=0.0, stddev=self.scale)
-                self._train_step(self.model, images, labels)
+                train_step(self.model, images, labels)
 
     def predict(self, x: np.ndarray, batch_size: int = 128, **kwargs) -> np.ndarray:  # type: ignore
         """

art/estimators/classification/tensorflow.py

@@ -800,6 +800,7 @@ class TensorFlowV2Classifier(ClassGradientsMixin, ClassifierMixin, TensorFlowV2E
         + [
             "input_shape",
             "loss_object",
+            "optimizer",
             "train_step",
         ]
     )
@@ -810,6 +811,7 @@ def __init__(
         nb_classes: int,
         input_shape: Tuple[int, ...],
         loss_object: Optional["tf.keras.losses.Loss"] = None,
+        optimizer: Optional["tf.keras.optimizers.Optimizer"] = None,
         train_step: Optional[Callable] = None,
         channels_first: bool = False,
         clip_values: Optional["CLIP_VALUES_TYPE"] = None,
@@ -824,10 +826,12 @@ def __init__(
         :param nb_classes: the number of classes in the classification task.
         :param input_shape: shape of one input for the classifier, e.g. for MNIST input_shape=(28, 28, 1).
         :param loss_object: The loss function for which to compute gradients. This parameter is applied for training
-            the model and computing gradients of the loss w.r.t. the input.
-        :type loss_object: `tf.keras.losses`
+               the model and computing gradients of the loss w.r.t. the input.
+        :param optimizer: The optimizer used to train the classifier.
         :param train_step: A function that applies a gradient update to the trainable variables with signature
-                           train_step(model, images, labels).
+               `train_step(model, images, labels)`. This will override the default training loop that uses the
+               provided `loss_object` and `optimizer` parameters. It is recommended to use the `@tf.function`
+               decorator, if possible, for efficient training.
         :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
@@ -853,6 +857,7 @@ def __init__(
         self.nb_classes = nb_classes
         self._input_shape = input_shape
         self._loss_object = loss_object
+        self._optimizer = optimizer
         self._train_step = train_step
 
         # Check if the loss function requires as input index labels instead of one-hot-encoded labels
@@ -879,6 +884,15 @@ def loss_object(self) -> "tf.keras.losses.Loss":
         """
         return self._loss_object  # type: ignore
 
+    @property
+    def optimizer(self) -> "tf.keras.optimizers.Optimizer":
+        """
+        Return the optimizer.
+
+        :return: The optimizer.
+        """
+        return self._optimizer  # type: ignore
+
     @property
     def train_step(self) -> Callable:
         """
@@ -949,9 +963,27 @@ def fit(self, x: np.ndarray, y: np.ndarray, batch_size: int = 128, nb_epochs: in
         import tensorflow as tf
 
         if self._train_step is None:  # pragma: no cover
-            raise TypeError(
-                "The training function `train_step` is required for fitting a model but it has not been " "defined."
-            )
+            if self._loss_object is None:  # pragma: no cover
+                raise TypeError(
+                    "A loss function `loss_object` or training function `train_step` is required for fitting the "
+                    "model, but it has not been defined."
+                )
+            if self._optimizer is None:  # pragma: no cover
+                raise ValueError(
+                    "An optimizer `optimizer` or training function `train_step` is required for fitting the "
+                    "model, but it has not been defined."
+                )
+
+            @tf.function
+            def train_step(model, images, labels):
+                with tf.GradientTape() as tape:
+                    predictions = model(images, training=True)
+                    loss = self.loss_object(labels, predictions)
+                gradients = tape.gradient(loss, model.trainable_variables)
+                self.optimizer.apply_gradients(zip(gradients, model.trainable_variables))
+
+        else:
+            train_step = self._train_step
 
         y = check_and_transform_label_format(y, nb_classes=self.nb_classes)
 
@@ -966,7 +998,7 @@ def fit(self, x: np.ndarray, y: np.ndarray, batch_size: int = 128, nb_epochs: in
 
         for _ in range(nb_epochs):
             for images, labels in train_ds:
-                self._train_step(self.model, images, labels)
+                train_step(self.model, images, labels)
 
     def fit_generator(self, generator: "DataGenerator", nb_epochs: int = 20, **kwargs) -> None:
         """
@@ -982,9 +1014,27 @@ def fit_generator(self, generator: "DataGenerator", nb_epochs: int = 20, **kwarg
         from art.data_generators import TensorFlowV2DataGenerator
 
         if self._train_step is None:  # pragma: no cover
-            raise TypeError(
-                "The training function `train_step` is required for fitting a model but it has not been " "defined."
-            )
+            if self._loss_object is None:  # pragma: no cover
+                raise TypeError(
+                    "A loss function `loss_object` or training function `train_step` is required for fitting the "
+                    "model, but it has not been defined."
+                )
+            if self._optimizer is None:  # pragma: no cover
+                raise ValueError(
+                    "An optimizer `optimizer` or training function `train_step` is required for fitting the "
+                    "model, but it has not been defined."
+                )
+
+            @tf.function
+            def train_step(model, images, labels):
+                with tf.GradientTape() as tape:
+                    predictions = model(images, training=True)
+                    loss = self.loss_object(labels, predictions)
+                gradients = tape.gradient(loss, model.trainable_variables)
+                self.optimizer.apply_gradients(zip(gradients, model.trainable_variables))
+
+        else:
+            train_step = self._train_step
 
         # Train directly in TensorFlow
         from art.preprocessing.standardisation_mean_std.tensorflow import StandardisationMeanStdTensorFlow
@@ -1004,7 +1054,7 @@ def fit_generator(self, generator: "DataGenerator", nb_epochs: int = 20, **kwarg
                 for i_batch, o_batch in generator.iterator:
                     if self._reduce_labels:
                         o_batch = tf.math.argmax(o_batch, axis=1)
-                    self._train_step(self._model, i_batch, o_batch)
+                    train_step(self._model, i_batch, o_batch)
         else:
             # Fit a generic data generator through the API
             super().fit_generator(generator, nb_epochs=nb_epochs)
@@ -1263,6 +1313,7 @@ def clone_for_refitting(
         clone._train_step = self._train_step  # pylint: disable=W0212
         clone._reduce_labels = self._reduce_labels  # pylint: disable=W0212
         clone._loss_object = self._loss_object  # pylint: disable=W0212
+        clone._optimizer = self._optimizer  # pylint: disable=W0212
         return clone
 
     def reset(self) -> None:
@@ -1401,8 +1452,8 @@ def save(self, filename: str, path: Optional[str] = None) -> None:
     def __repr__(self):
         repr_ = (
             f"{self.__module__ + '.' + self.__class__.__name__}(model={self._model}, nb_classes={self.nb_classes}, "
-            f"input_shape={self._input_shape}, loss_object={self._loss_object}, train_step={self._train_step}, "
-            f"channels_first={self.channels_first}, clip_values={self.clip_values!r}, "
+            f"input_shape={self._input_shape}, loss_object={self._loss_object}, optimizer={self.optimizer}, "
+            f"train_step={self._train_step}, channels_first={self.channels_first}, clip_values={self.clip_values!r}, "
             f"preprocessing_defences={self.preprocessing_defences}, "
             f"postprocessing_defences={self.postprocessing_defences}, preprocessing={self.preprocessing})"
         )

examples/get_started_tensorflow_v2.py

@@ -52,26 +52,16 @@ def call(self, x):
         return x
 
 
-optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)
-
-
-def train_step(model, images, labels):
-    with tf.GradientTape() as tape:
-        predictions = model(images, training=True)
-        loss = loss_object(labels, predictions)
-    gradients = tape.gradient(loss, model.trainable_variables)
-    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
-
-
 model = TensorFlowModel()
 loss_object = tf.keras.losses.CategoricalCrossentropy(from_logits=True)
+optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)
 
 # Step 3: Create the ART classifier
 
 classifier = TensorFlowV2Classifier(
     model=model,
     loss_object=loss_object,
-    train_step=train_step,
+    optimizer=optimizer,
     nb_classes=10,
     input_shape=(28, 28, 1),
     clip_values=(0, 1),