add learning rate schedulers to tensorflow estimators

Signed-off-by: Farhan Ahmed <Farhan.Ahmed@ibm.com>

art/estimators/certification/randomized_smoothing/macer/pytorch.py

@@ -177,7 +177,9 @@ def fit(  # pylint: disable=W0221
         dataset = TensorDataset(x_tensor, y_tensor)
         dataloader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True, drop_last=drop_last)
 
-        m = torch.distributions.normal.Normal(torch.tensor([0.0], device=self.device), torch.tensor([1.0], device=self.device))
+        m = torch.distributions.normal.Normal(
+            torch.tensor([0.0], device=self.device), torch.tensor([1.0], device=self.device)
+        )
 
         # Start training
         for _ in tqdm(range(nb_epochs)):

art/estimators/certification/randomized_smoothing/tensorflow.py

@@ -137,8 +137,9 @@ def fit(self, x: np.ndarray, y: np.ndarray, batch_size: int = 128, nb_epochs: in
                   shape (nb_samples,).
         :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
-               TensorFlow and providing it takes no effect.
+        :param kwargs: Dictionary of framework-specific arguments. This parameter currently only supports
+                       "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
 
@@ -165,6 +166,8 @@ def train_step(model, images, labels):
         else:
             train_step = self._train_step
 
+        scheduler = kwargs.get("scheduler")
+
         y = check_and_transform_label_format(y, nb_classes=self.nb_classes)
 
         # Apply preprocessing
@@ -176,12 +179,15 @@ def train_step(model, images, labels):
 
         train_ds = tf.data.Dataset.from_tensor_slices((x_preprocessed, y_preprocessed)).shuffle(10000).batch(batch_size)
 
-        for _ in tqdm(range(nb_epochs)):
+        for epoch in tqdm(range(nb_epochs)):
             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)
                 train_step(self.model, images, labels)
 
+            if scheduler is not None:
+                scheduler(epoch)
+
     def predict(self, x: np.ndarray, batch_size: int = 128, **kwargs) -> np.ndarray:  # type: ignore
         """
         Perform prediction of the given classifier for a batch of inputs, taking an expectation over transformations.

art/estimators/classification/tensorflow.py

@@ -957,8 +957,9 @@ def fit(self, x: np.ndarray, y: np.ndarray, batch_size: int = 128, nb_epochs: in
                   shape (nb_samples,).
         :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
-               TensorFlow and providing it takes no effect.
+        :param kwargs: Dictionary of framework-specific arguments. This parameter currently only supports
+                       "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
 
@@ -985,6 +986,8 @@ def train_step(model, images, labels):
         else:
             train_step = self._train_step
 
+        scheduler = kwargs.get("scheduler")
+
         y = check_and_transform_label_format(y, nb_classes=self.nb_classes)
 
         # Apply preprocessing
@@ -996,19 +999,23 @@ def train_step(model, images, labels):
 
         train_ds = tf.data.Dataset.from_tensor_slices((x_preprocessed, y_preprocessed)).shuffle(10000).batch(batch_size)
 
-        for _ in range(nb_epochs):
+        for epoch in range(nb_epochs):
             for images, labels in train_ds:
                 train_step(self.model, images, labels)
 
+            if scheduler is not None:
+                scheduler(epoch)
+
     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. If the generator can be used for native
                           training in TensorFlow, it will.
         :param nb_epochs: Number of epochs to use for training.
-        :param kwargs: Dictionary of framework-specific arguments. This parameter is not currently supported for
-               TensorFlow and providing it takes no effect.
+        :param kwargs: Dictionary of framework-specific arguments. This parameter currently only supports
+                       "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
         from art.data_generators import TensorFlowV2DataGenerator
@@ -1036,6 +1043,8 @@ def train_step(model, images, labels):
         else:
             train_step = self._train_step
 
+        scheduler = kwargs.get("scheduler")
+
         # Train directly in TensorFlow
         from art.preprocessing.standardisation_mean_std.tensorflow import StandardisationMeanStdTensorFlow
 
@@ -1050,11 +1059,14 @@ def train_step(model, images, labels):
                 == (0, 1)
             )
         ):
-            for _ in range(nb_epochs):
+            for epoch in range(nb_epochs):
                 for i_batch, o_batch in generator.iterator:
                     if self._reduce_labels:
                         o_batch = tf.math.argmax(o_batch, axis=1)
                     train_step(self._model, i_batch, o_batch)
+
+                if scheduler is not None:
+                    scheduler(epoch)
         else:
             # Fit a generic data generator through the API
             super().fit_generator(generator, nb_epochs=nb_epochs)