Support for many new model output types (#93)

* General model wrappers and methods supporting multi-label classifiers
* Support for pytorch multi-label binary classifier
* New model output types + single implementation of score method that supports multiple output types. 
* Anonymization with pytorch multi-output binary model
* Support for multi-label binary models in minimizer. 
* Support for multi-label logits/probabilities
---------
Signed-off-by: abigailt <abigailt@il.ibm.com>

apt/utils/models/pytorch_model.py

@@ -3,17 +3,22 @@ import os
 import shutil
 import logging
 
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Union, List, TYPE_CHECKING
 import numpy as np
 import torch
 from torch.utils.data import DataLoader, TensorDataset
 
 from art.utils import check_and_transform_label_format
-from apt.utils.datasets.datasets import PytorchData
-from apt.utils.models import Model, ModelOutputType
-from apt.utils.datasets import OUTPUT_DATA_ARRAY_TYPE
+from apt.utils.datasets.datasets import PytorchData, DatasetWithPredictions, ArrayDataset
+from apt.utils.models import Model, ModelOutputType, is_multi_label, is_multi_label_binary, is_binary
+from apt.utils.datasets import OUTPUT_DATA_ARRAY_TYPE, array2numpy
 from art.estimators.classification.pytorch import PyTorchClassifier as ArtPyTorchClassifier
 
+if TYPE_CHECKING:
+    from art.utils import CLIP_VALUES_TYPE, PREPROCESSING_TYPE
+    from art.defences.preprocessor import Preprocessor
+    from art.defences.postprocessor import Postprocessor
+
 
 logger = logging.getLogger(__name__)
 
@@ -30,16 +35,46 @@ class PyTorchClassifierWrapper(ArtPyTorchClassifier):
     Extension for Pytorch ART model
     """
 
+    def __init__(
+            self,
+            model: "torch.nn.Module",
+            loss: "torch.nn.modules.loss._Loss",
+            input_shape: Tuple[int, ...],
+            nb_classes: int,
+            output_type: ModelOutputType,
+            optimizer: Optional["torch.optim.Optimizer"] = None,  # type: ignore
+            use_amp: bool = False,
+            opt_level: str = "O1",
+            loss_scale: Optional[Union[float, str]] = "dynamic",
+            channels_first: bool = True,
+            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),
+            device_type: str = "gpu",
+    ):
+        super().__init__(model, loss, input_shape, nb_classes, optimizer, use_amp, opt_level, loss_scale,
+                         channels_first, clip_values, preprocessing_defences, postprocessing_defences, preprocessing,
+                         device_type)
+        self._is_single_binary = not is_multi_label(output_type) and is_binary(output_type)
+        self._is_multi_label = is_multi_label(output_type)
+        self._is_multi_label_binary = is_multi_label_binary(output_type)
+
     def get_step_correct(self, outputs, targets) -> int:
         """
         Get number of correctly classified labels.
         """
+        # here everything is torch tensors
         if len(outputs) != len(targets):
             raise ValueError("outputs and targets should be the same length.")
-        if self.nb_classes > 1:
-            return int(torch.sum(torch.argmax(outputs, axis=-1) == targets).item())
+        if self._is_single_binary:
+            return int(torch.sum(torch.round(outputs) == targets).item())
+        elif self._is_multi_label:
+            if self._is_multi_label_binary:
+                outputs = torch.round(outputs)
+            return int(torch.sum(targets == outputs).item())
         else:
-            return int(torch.sum(torch.round(outputs, axis=-1) == targets).item())
+            return int(torch.sum(torch.argmax(outputs, axis=-1) == targets).item())
 
     def _eval(self, loader: DataLoader):
         """
@@ -93,6 +128,7 @@ class PyTorchClassifierWrapper(ArtPyTorchClassifier):
         :param kwargs: Dictionary of framework-specific arguments. This parameter is not currently
                        supported for PyTorch and providing it takes no effect.
         """
+
         # Put the model in the training mode
         self._model.train()
 
@@ -156,6 +192,61 @@ class PyTorchClassifierWrapper(ArtPyTorchClassifier):
                 else:
                     self.save_checkpoint_state_dict(is_best=best_acc <= val_acc, path=path)
 
+    def predict(
+            self, x: np.ndarray, batch_size: int = 128, training_mode: bool = False, **kwargs
+    ) -> np.ndarray:
+        """
+        Perform prediction for a batch of inputs.
+
+        :param x: Input samples.
+        :param batch_size: Size of batches.
+        :param training_mode: `True` for model set to training mode and `'False` for model set to evaluation mode.
+        :return: Array of predictions of shape `(nb_inputs, nb_classes)`.
+        """
+        import torch
+
+        # Set model mode
+        self._model.train(mode=training_mode)
+
+        # 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]),
+            )
+
+            with torch.no_grad():
+                model_outputs = self._model(torch.from_numpy(x_preprocessed[begin:end]).to(self._device))
+            output = model_outputs[-1]
+
+            if isinstance(output, tuple):
+                output_list = []
+                for o in output:
+                    o = o.detach().cpu().numpy().astype(np.float32)
+                    output_list.append(o)
+                output_np = np.array(output_list)
+                output_np = np.swapaxes(output_np, 0, 1)
+                results_list.append(output_np)
+            else:
+                output = output.detach().cpu().numpy().astype(np.float32)
+                if len(output.shape) == 1:
+                    output = np.expand_dims(output, axis=1).astype(np.float32)
+                results_list.append(output)
+
+        results = np.vstack(results_list)
+
+        # Apply postprocessing
+        predictions = self._apply_postprocessing(preds=results, fit=False)
+
+        return predictions
+
     def save_checkpoint_state_dict(self, is_best: bool, path=os.getcwd(), filename="latest.tar") -> None:
         """
         Saves checkpoint as latest.tar or best.tar.
@@ -319,7 +410,8 @@ class PyTorchClassifier(PyTorchModel):
         super().__init__(model, output_type, black_box_access, unlimited_queries, **kwargs)
         self._loss = loss
         self._optimizer = optimizer
-        self._art_model = PyTorchClassifierWrapper(model, loss, input_shape, nb_classes, optimizer)
+        self._nb_classes = nb_classes
+        self._art_model = PyTorchClassifierWrapper(model, loss, input_shape, nb_classes, output_type, optimizer)
 
     @property
     def loss(self):
@@ -398,7 +490,7 @@ class PyTorchClassifier(PyTorchModel):
         :type x: `np.ndarray` or `pandas.DataFrame`
         :return: Predictions from the model (class probabilities, if supported).
         """
-        return self._art_model.predict(x.get_samples(), **kwargs)
+        return array2numpy(self._art_model.predict(x.get_samples(), **kwargs))
 
     def score(self, test_data: PytorchData, **kwargs):
         """
@@ -406,18 +498,20 @@ class PyTorchClassifier(PyTorchModel):
 
         :param test_data: Test data.
         :type test_data: `PytorchData`
+        :param binary_threshold: The threshold to use on binary classification probabilities to assign the positive
+                                 class.
+        :type binary_threshold: float, optional. Default is 0.5.
+        :param apply_non_linearity: A non-linear function to apply to the result of the 'predict' method, in case the
+                                    model outputs logits (e.g., sigmoid).
+        :type apply_non_linearity: Callable, should be possible to apply directly to the numpy output of the 'predict'
+                                   method, optional.
         :return: the score as float (between 0 and 1)
         """
-        y = test_data.get_labels()
+        # numpy arrays
         predicted = self.predict(test_data)
-        # binary classification, single column of probabilities
-        if self._art_model.nb_classes == 2 and (len(predicted.shape) == 1 or predicted.shape[1] == 1):
-            if len(predicted.shape) > 1:
-                y = check_and_transform_label_format(y, self._art_model.nb_classes, return_one_hot=False)
-            return np.count_nonzero(y == (predicted > 0.5)) / predicted.shape[0]
-        else:
-            y = check_and_transform_label_format(y, self._art_model.nb_classes)
-            return np.count_nonzero(np.argmax(y, axis=1) == np.argmax(predicted, axis=1)) / predicted.shape[0]
+        kwargs['predictions'] = DatasetWithPredictions(pred=predicted)
+        kwargs['nb_classes'] = self._nb_classes
+        return super().score(ArrayDataset(test_data.get_samples(), test_data.get_labels()), **kwargs)
 
     def load_checkpoint_state_dict_by_path(self, model_name: str, path: str = None):
         """