progress bar development

Signed-off-by: GiulioZizzo <[email protected]>

Author: GiulioZizzo

.github/workflows/ci-pytorch.yml

@@ -16,6 +16,7 @@ on:
     branches:
       - main
       - dev*
+      - hf_notebook_updates
 
   # Run scheduled CI flow daily
   schedule:

.github/workflows/ci-style-checks.yml

@@ -16,6 +16,7 @@ on:
     branches:
       - main
       - dev*
+      - hf_notebook_updates
 
   # Run scheduled CI flow daily
   schedule:

art/estimators/classification/pytorch.py

@@ -26,6 +26,7 @@
 import os
 import time
 from typing import Any, Dict, List, Optional, Tuple, Union, TYPE_CHECKING
+from tqdm.auto import tqdm
 
 import numpy as np
 import six
@@ -389,12 +390,14 @@ def fit(  # pylint: disable=W0221
                           the batch size. If ``False`` and the size of dataset is not divisible by the batch size, then
                           the last batch will be smaller. (default: ``False``)
         :param scheduler: Learning rate scheduler to run at the start of every epoch.
-        :param kwargs: Dictionary of framework-specific arguments. This parameter is not currently supported for PyTorch
-                       and providing it takes no effect.
+        :param kwargs: Dictionary of framework-specific arguments. Currently supports "display_progress_bar" to
+                       display training progress.
         """
         import torch
         from torch.utils.data import TensorDataset, DataLoader
 
+        display_progress_bar = kwargs.get("display_progress_bar", False)
+
         # Set model mode
         self._model.train(mode=training_mode)
 
@@ -416,8 +419,8 @@ def fit(  # pylint: disable=W0221
         dataloader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True, drop_last=drop_last)
 
         # Start training
-        for _ in range(nb_epochs):
-            for x_batch, y_batch in dataloader:
+        for _ in tqdm(range(nb_epochs), disable=not display_progress_bar, desc="Epochs"):
+            for x_batch, y_batch in tqdm(dataloader, disable=not display_progress_bar, desc="Batches"):
                 # Move inputs to device
                 x_batch = x_batch.to(self._device)
                 y_batch = y_batch.to(self._device)
@@ -459,12 +462,14 @@ def fit_generator(self, generator: "DataGenerator", nb_epochs: int = 20, **kwarg
 
         :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. This parameter is not currently supported for PyTorch
-               and providing it takes no effect.
+        :param kwargs: Dictionary of framework-specific arguments. Currently supports "display_progress_bar" to
+                       display training progress.
         """
         import torch
         from art.data_generators import PyTorchDataGenerator
 
+        display_progress_bar = kwargs.get("display_progress_bar", False)
+
         # Put the model in the training mode
         self._model.train()
 
@@ -485,8 +490,8 @@ def fit_generator(self, generator: "DataGenerator", nb_epochs: int = 20, **kwarg
                 == (0, 1)
             )
         ):
-            for _ in range(nb_epochs):
-                for i_batch, o_batch in generator.iterator:
+            for _ in tqdm(range(nb_epochs), disable=not display_progress_bar, desc="Epochs"):
+                for i_batch, o_batch in tqdm(generator.iterator, disable=not display_progress_bar, desc="Batches"):
                     if isinstance(i_batch, np.ndarray):
                         i_batch = torch.from_numpy(i_batch).to(self._device)
                     else:
@@ -495,7 +500,10 @@ def fit_generator(self, generator: "DataGenerator", nb_epochs: int = 20, **kwarg
                     if isinstance(o_batch, np.ndarray):
                         o_batch = torch.argmax(torch.from_numpy(o_batch).to(self._device), dim=1)
                     else:
-                        o_batch = torch.argmax(o_batch.to(self._device), dim=1)
+                        if o_batch.dim() > 1:
+                            o_batch = torch.argmax(o_batch.to(self._device), dim=1)
+                        else:
+                            o_batch = o_batch.to(self._device)
 
                     # Zero the parameter gradients
                     self._optimizer.zero_grad()

art/estimators/classification/tensorflow.py

@@ -27,6 +27,7 @@
 import shutil
 import time
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union, TYPE_CHECKING
+from tqdm.auto import tqdm
 
 import numpy as np
 import six
@@ -957,12 +958,14 @@ 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 currently only supports
+        :param kwargs: Dictionary of framework-specific arguments. This parameter currently supports
                        "scheduler" which is an optional function that will be called at the end of every
-                       epoch to adjust the learning rate.
+                       epoch to adjust the learning rate, and "display_progress_bar" to display training progress.
         """
         import tensorflow as tf
 
+        display_progress_bar = kwargs.get("display_progress_bar", False)
+
         if self._train_step is None:  # pragma: no cover
             if self._loss_object is None:  # pragma: no cover
                 raise TypeError(
@@ -999,8 +1002,8 @@ def train_step(model, images, labels):
 
         train_ds = tf.data.Dataset.from_tensor_slices((x_preprocessed, y_preprocessed)).shuffle(10000).batch(batch_size)
 
-        for epoch in range(nb_epochs):
-            for images, labels in train_ds:
+        for epoch in tqdm(range(nb_epochs), disable=not display_progress_bar, desc="Epochs"):
+            for images, labels in tqdm(train_ds, disable=not display_progress_bar, desc="Batches"):
                 train_step(self.model, images, labels)
 
             if scheduler is not None:
@@ -1013,13 +1016,15 @@ def fit_generator(self, generator: "DataGenerator", nb_epochs: int = 20, **kwarg
         :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 currently only supports
+        :param kwargs: Dictionary of framework-specific arguments. This parameter currently supports
                        "scheduler" which is an optional function that will be called at the end of every
-                       epoch to adjust the learning rate.
+                       epoch to adjust the learning rate, and "display_progress_bar" to display training progress.
         """
         import tensorflow as tf
         from art.data_generators import TensorFlowV2DataGenerator
 
+        display_progress_bar = kwargs.get("display_progress_bar", False)
+
         if self._train_step is None:  # pragma: no cover
             if self._loss_object is None:  # pragma: no cover
                 raise TypeError(
@@ -1059,8 +1064,8 @@ def train_step(model, images, labels):
                 == (0, 1)
             )
         ):
-            for epoch in range(nb_epochs):
-                for i_batch, o_batch in generator.iterator:
+            for epoch in tqdm(range(nb_epochs), disable=not display_progress_bar, desc="Epochs"):
+                for i_batch, o_batch in tqdm(generator.iterator, disable=not display_progress_bar, desc="Batches"):
                     if self._reduce_labels:
                         o_batch = tf.math.argmax(o_batch, axis=1)
                     train_step(self._model, i_batch, o_batch)