Improve prediction performance

Signed-off-by: Beat Buesser <[email protected]>

Author: beat-buesser

art/estimators/classification/keras.py

@@ -97,6 +97,8 @@ def __init__(
                              layer with this index will be considered for computing gradients. For models with only one
                              output layer this values is not required.
         """
+        import tensorflow as tf
+
         super().__init__(
             model=model,
             clip_values=clip_values,
@@ -130,6 +132,12 @@ def __init__(
         self._input_shape = tuple(self._input.shape[1:])
         self._layer_names = self._get_layers()
 
+        @tf.function(reduce_retracing=True)  # Compile this for speed
+        def _forward_pass(model, x, training, batch_size):
+            return model(x, training=training, batch_size=batch_size, verbose=False)
+
+        self._forward_pass = _forward_pass
+
     @property
     def input_shape(self) -> tuple[int, ...]:
         """
@@ -397,15 +405,14 @@ def predict(self, x: np.ndarray, batch_size: int = 128, training_mode: bool = Fa
         x_preprocessed, _ = self._apply_preprocessing(x, y=None, fit=False)
 
         # Run predictions with batching
-        if training_mode:
-            predictions = self._model(x_preprocessed, training=training_mode, verbose=False)
-        else:
-            predictions = self._model.predict(x_preprocessed, batch_size=batch_size, verbose=False)
+        predictions = self._forward_pass(
+            self._model, x_preprocessed, training=training_mode, batch_size=batch_size
+        )  # Fast, compiled call
 
         # Apply postprocessing
-        predictions = self._apply_postprocessing(preds=predictions, fit=False)
+        predictions_post = self._apply_postprocessing(preds=predictions.numpy(), fit=False)
 
-        return predictions
+        return predictions_post
 
     def fit(
         self, x: np.ndarray, y: np.ndarray, batch_size: int = 128, nb_epochs: int = 20, verbose: bool = False, **kwargs

art/estimators/classification/tensorflow.py

@@ -120,6 +120,12 @@ def __init__(
         else:
             self._reduce_labels = False
 
+        @tf.function(reduce_retracing=True)  # Compile this for speed
+        def _forward_pass(model, x, training):
+            return model(x, training=training)
+
+        self._forward_pass = _forward_pass
+
     @property
     def input_shape(self) -> tuple[int, ...]:
         """
@@ -168,24 +174,13 @@ def predict(self, x: np.ndarray, batch_size: int = 128, training_mode: bool = Fa
         # Apply preprocessing
         x_preprocessed, _ = self._apply_preprocessing(x, y=None, fit=False)
 
-        # Run prediction with batch processing
-        results_list = []
-        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]),
-            )
-
-            # Run prediction
-            results_list.append(self._model(x_preprocessed[begin:end], training=training_mode))
-
-        results = np.vstack(results_list)
+        # Run predictions with batching
+        predictions = self._forward_pass(self._model, x_preprocessed, training=training_mode)  # Fast, compiled call
 
         # Apply postprocessing
-        predictions = self._apply_postprocessing(preds=results, fit=False)
-        return predictions
+        predictions_post = self._apply_postprocessing(preds=predictions.numpy(), fit=False)
+
+        return predictions_post
 
     def _predict_framework(self, x: "tf.Tensor", training_mode: bool = False) -> "tf.Tensor":
         """

conftest.py

@@ -150,7 +150,7 @@ def _image_dl_estimator_for_attack(attack, defended=False, **kwargs):
             potential_classifier, _ = image_dl_estimator_defended(**kwargs)
         else:
             potential_classifier, _ = image_dl_estimator(**kwargs)
-        image_dl_estimator_for_attack
+
         classifier_list = [potential_classifier]
         classifier_tested = [
             potential_classifier
@@ -459,16 +459,19 @@ def _image_dl_estimator(functional=False, **kwargs):
                             image_dl_estimator.__name__,
                             framework,
                         )
+
         if framework == "tensorflow2":
             if wildcard is False and functional is False:
                 classifier, sess = get_image_classifier_tf(**kwargs, framework=framework)
                 return classifier, sess
+
         if framework == "pytorch":
             if not wildcard:
                 if functional:
                     classifier = get_image_classifier_pt_functional(**kwargs)
                 else:
                     classifier = get_image_classifier_pt(**kwargs)
+
         if framework == "kerastf":
             if wildcard:
                 classifier = get_image_classifier_kr_tf_with_wildcard(**kwargs)

tests/attacks/evasion/test_sign_opt.py

@@ -54,7 +54,7 @@ def fix_get_mnist_subset_large(get_mnist_dataset):
 def test_tabular(art_warning, tabular_dl_estimator, framework, get_iris_dataset, clipped_classifier, targeted):
     try:
         classifier = tabular_dl_estimator(clipped=clipped_classifier)
-        attack = SignOPTAttack(classifier, targeted=targeted, max_iter=1000, query_limit=4000, verbose=False)
+        attack = SignOPTAttack(classifier, targeted=targeted, num_trial=10, max_iter=100, query_limit=40, verbose=True)
         if targeted:
             backend_targeted_tabular(attack, get_iris_dataset)
         else:

tests/utils.py

@@ -665,14 +665,15 @@ def get_image_classifier_kr_tf(loss_name="categorical_crossentropy", loss_type="
     import tensorflow as tf
 
     # pylint: disable=E0401
-    from tensorflow.keras.layers import Conv2D, Dense, Flatten, MaxPooling2D
+    from tensorflow.keras.layers import Conv2D, Dense, Flatten, MaxPooling2D, Input
     from tensorflow.keras.models import Sequential
 
     from art.estimators.classification.keras import KerasClassifier
 
     # Create simple CNN
     model = Sequential()
-    model.add(Conv2D(1, kernel_size=(7, 7), activation="relu", input_shape=(28, 28, 1)))
+    model.add(Input((28, 28, 1)))
+    model.add(Conv2D(1, kernel_size=(7, 7), activation="relu"))
     model.layers[-1].set_weights(
         [_kr_tf_weights_loader("MNIST", "W", "CONV2D"), _kr_tf_weights_loader("MNIST", "B", "CONV2D")]
     )