Craft updates to prepare Cockatiel.

fredericboisnard · fredericboisnard · commit 5743025eabb9 · 2023-11-27T16:56:14.000+01:00
This patch separates the methods dedicated to image
processing/plotting into different sub classes, so that
Cockatiel and NLP related classes will not have to inheritate
these.

Signed-off-by: Frederic Boisnard &lt;frederic.boisnard@irt-saintexupery.com&gt;
diff --git a/xplique/concepts/craft.py b/xplique/concepts/craft.py
@@ -129,7 +129,7 @@ def set_concept_attribution_cmap(self, cmaps: Optional[Union[Tuple, str]]=None):
         if len(self.cmaps) < len(self.most_important_concepts):
             raise RuntimeError(f'Not enough colors in cmaps ({len(self.cmaps)}) ' \
                                f'compared to the number of important concepts ' \
-                               '({len(self.most_important_concepts)})')
+                               f'({len(self.most_important_concepts)})')
 
 class DisplayImportancesOrder(Enum):
     """
@@ -455,6 +455,44 @@ def _show(img, **kwargs):
         plt.imshow(img, **kwargs)
         plt.axis('off')
 
+    def _gen_best_concepts_crops(self,
+                                 nb_crops: int = 10,
+                                 nb_most_important_concepts: int = None) \
+                                 -> Tuple[int, float, np.ndarray]:
+        """
+        Generate the best concept crops for each concept.
+
+        Parameters
+        ----------
+        nb_crops : int
+            The number of crops (patches) to display per concept. Defaults to 10.
+        nb_most_important_concepts : int
+            The number of concepts to consider. If provided, only take into account
+            nb_most_important_concepts, otherwise use them all.
+            Default is None.
+        Returns
+        -------
+        Tuple
+            A tuple containing:
+            - The current concept id.
+            - The overall importance score for this concept.
+            - An array containing the best crops for this concept.
+        """
+        most_important_concepts = self.sensitivity.most_important_concepts
+        if nb_most_important_concepts is not None:
+            most_important_concepts = most_important_concepts[:nb_most_important_concepts]
+
+        for c_id in most_important_concepts:
+            best_crops_ids = np.argsort(self.factorization.crops_u[:, c_id])[::-1][:nb_crops]
+            best_crops = np.array(self.factorization.crops)[best_crops_ids]
+            c_id_importance = self.sensitivity.importances[c_id]
+            yield c_id, c_id_importance, best_crops
+
+
+class CraftImageVisualizationMixin():
+    """
+    Class containing image visualization methods for Craft.
+    """
     def plot_concepts_crops(self,
                             nb_crops: int = 10,
                             nb_most_important_concepts: int = None,
@@ -474,17 +512,11 @@ def plot_concepts_crops(self,
             If True, then print the importance value of each concept,
             otherwise no textual output will be printed.
         """
-        most_important_concepts = self.sensitivity.most_important_concepts
-        if nb_most_important_concepts is not None:
-            most_important_concepts = most_important_concepts[:nb_most_important_concepts]
-
-        for c_id in most_important_concepts:
-            best_crops_ids = np.argsort(self.factorization.crops_u[:, c_id])[::-1][:nb_crops]
-            best_crops = np.array(self.factorization.crops)[best_crops_ids]
-
+        for c_id, c_id_importance, best_crops in \
+            self._gen_best_concepts_crops(nb_crops, nb_most_important_concepts):
             if verbose:
                 print(f"Concept {c_id} has an importance value of " \
-                    f"{self.sensitivity.importances[c_id]:.2f}")
+                        f"{c_id_importance:.2f}")
             plt.figure(figsize=(7, (2.5/2)*ceil(nb_crops/5)))
             for i in range(nb_crops):
                 plt.subplot(ceil(nb_crops/5), 5, i+1)
diff --git a/xplique/concepts/craft_manager.py b/xplique/concepts/craft_manager.py
@@ -104,6 +104,11 @@ def estimate_importance(self, nb_design: int = 32, verbose: bool = False):
                 print(f'Estimating importances for class {class_of_interest} ')
             craft_instance.estimate_importance(nb_design=nb_design)
 
+
+class CraftManagerImageVisualizationMixin():
+    """
+    Class containing image visualization methods for CraftManager.
+    """
     def plot_concepts_importances(self,
                                   class_id: int,
                                   nb_most_important_concepts: int = 5,
diff --git a/xplique/concepts/craft_tf.py b/xplique/concepts/craft_tf.py
@@ -7,13 +7,12 @@
 import tensorflow as tf
 import numpy as np
 
-from .craft import BaseCraft
-from .craft_manager import BaseCraftManager
-
-class CraftTf(BaseCraft):
+from .craft import BaseCraft, CraftImageVisualizationMixin
+from .craft_manager import BaseCraftManager, CraftManagerImageVisualizationMixin
 
+class BaseCraftTf(BaseCraft):
     """
-    Class implementing the CRAFT Concept Extraction Mechanism on Tensorflow.
+    Base class implementing the CRAFT Concept Extraction Mechanism on Tensorflow.
 
     Parameters
     ----------
@@ -137,10 +136,34 @@ def _to_np_array(self, inputs: tf.Tensor, dtype: type):
         """
         return np.array(inputs, dtype)
 
+class CraftTf(BaseCraftTf, CraftImageVisualizationMixin):
+    """
+    Base class implementing the CRAFT Concept Extraction Mechanism on Tensorflow,
+    adapted for image processing.
 
-class CraftManagerTf(BaseCraftManager):
+    Parameters
+    ----------
+    input_to_latent_model
+        The first part of the model taking an input and returning
+        positive activations, g(.) in the original paper.
+        Must be a Tensorflow model (tf.keras.engine.base_layer.Layer) accepting
+        data of shape (n_samples, height, width, channels).
+    latent_to_logit_model
+        The second part of the model taking activation and returning
+        logits, h(.) in the original paper.
+        Must be a Tensorflow model (tf.keras.engine.base_layer.Layer).
+    number_of_concepts
+        The number of concepts to extract. Default is 20.
+    batch_size
+        The batch size to use during training and prediction. Default is 64.
+    patch_size
+        The size of the patches to extract from the input data. Default is 64.
     """
-    Class implementing the CraftManager on Tensorflow.
+
+
+class BaseCraftManagerTf(BaseCraftManager):
+    """
+    Base class implementing the CraftManager on Tensorflow.
     This manager creates one CraftTf instance per class to explain.
 
     Parameters
@@ -198,3 +221,33 @@ def compute_predictions(self):
         y_preds = np.array(tf.argmax(self.latent_to_logit_model.predict(
                             self.input_to_latent_model.predict(self.inputs)), 1))
         return y_preds
+
+class CraftManagerTf(BaseCraftManagerTf, CraftManagerImageVisualizationMixin):
+    """
+    Class implementing the CraftManager on Tensorflow, adapted for image processing.
+    This manager creates one CraftTf instance per class to explain.
+
+    Parameters
+    ----------
+    input_to_latent_model
+        The first part of the model taking an input and returning
+        positive activations, g(.) in the original paper.
+        Must return positive activations.
+    latent_to_logit_model
+        The second part of the model taking activation and returning
+        logits, h(.) in the original paper.
+    inputs
+        Input data of shape (n_samples, height, width, channels).
+        (x1, x2, ..., xn) in the paper.
+    labels
+        Labels of the inputs of shape (n_samples, class_id)
+    list_of_class_of_interest
+        A list of the classes id to explain. The manager will instanciate one
+        CraftTf object per element of this list.
+    number_of_concepts
+        The number of concepts to extract. Default is 20.
+    batch_size
+        The batch size to use during training and prediction. Default is 64.
+    patch_size
+        The size of the patches (crops) to extract from the input data. Default is 64.
+    """
diff --git a/xplique/concepts/craft_torch.py b/xplique/concepts/craft_torch.py
@@ -3,13 +3,14 @@
 """
 
 from typing import Callable, Optional, Tuple
+from types import MethodType
 from math import ceil
 import torch
 from torch import nn
 import numpy as np
 
-from .craft import BaseCraft
-from .craft_manager import BaseCraftManager
+from .craft import BaseCraft, CraftImageVisualizationMixin
+from .craft_manager import BaseCraftManager, CraftManagerImageVisualizationMixin
 
 def _batch_inference(model: torch.nn.Module,
                      dataset: torch.Tensor,
@@ -59,9 +60,9 @@ def _batch_inference(model: torch.nn.Module,
     return results
 
 
-class CraftTorch(BaseCraft):
+class BaseCraftTorch(BaseCraft):
     """
-    Class Implementing the CRAFT Concept Extraction Mechanism on Pytorch.
+    Base class implementing the CRAFT Concept Extraction Mechanism on Pytorch.
 
     Parameters
     ----------
@@ -83,7 +84,6 @@ class CraftTorch(BaseCraft):
     device
         The device to use. Default is 'cuda'.
     """
-
     def __init__(self, input_to_latent_model: Callable,
                        latent_to_logit_model: Callable,
                        number_of_concepts: int = 20,
@@ -96,11 +96,13 @@ def __init__(self, input_to_latent_model: Callable,
         self.device = device
 
         # Check model type
+        is_method = isinstance(input_to_latent_model, MethodType) & \
+                    isinstance(latent_to_logit_model, MethodType)
         is_torch_model = issubclass(type(input_to_latent_model), torch.nn.modules.module.Module) & \
                          issubclass(type(latent_to_logit_model), torch.nn.modules.module.Module)
-        if not is_torch_model:
+        if not (is_method or is_torch_model):
             raise TypeError('input_to_latent_model and latent_to_logit_model are not ' \
-                            'Pytorch modules')
+                            'Pytorch modules nor methods')
 
     def _latent_predict(self, inputs: torch.Tensor, resize=None) -> torch.Tensor:
         """
@@ -204,10 +206,35 @@ def _to_np_array(self, inputs: torch.Tensor, dtype: type = None):
             return res.astype(dtype)
         return res
 
+class CraftTorch(BaseCraftTorch, CraftImageVisualizationMixin):
+    """
+    Class Implementing the CRAFT Concept Extraction Mechanism on Pytorch,
+    adpated for image processing.
 
-class CraftManagerTorch(BaseCraftManager):
+    Parameters
+    ----------
+    input_to_latent_model
+        The first part of the model taking an input and returning
+        positive activations, g(.) in the original paper.
+        Must be a Pytorch model (torch.nn.modules.module.Module) accepting
+        data of shape (n_samples, channels, height, width).
+    latent_to_logit_model
+        The second part of the model taking activation and returning
+        logits, h(.) in the original paper.
+        Must be a Pytorch model (torch.nn.modules.module.Module).
+    number_of_concepts
+        The number of concepts to extract. Default is 20.
+    batch_size
+        The batch size to use during training and prediction. Default is 64.
+    patch_size
+        The size of the patches (crops) to extract from the input data. Default is 64.
+    device
+        The device to use. Default is 'cuda'.
+    """
+
+class BaseCraftManagerTorch(BaseCraftManager):
     """
-    Class implementing the CraftManager on Tensorflow.
+    Base class implementing the CraftManager on Pytorch.
     This manager creates one CraftTorch instance per class to explain.
 
     Parameters
@@ -270,3 +297,33 @@ def compute_predictions(self):
                                        device=self.device)
         y_preds = np.array(torch.argmax(activations, -1))  # pylint disable=no-member
         return y_preds
+
+class CraftManagerTorch(BaseCraftManagerTorch, CraftManagerImageVisualizationMixin):
+   """
+    Class implementing the CraftManager on Pytorch, adapted for image processing.
+    This manager creates one CraftTorch instance per class to explain.
+
+    Parameters
+    ----------
+    input_to_latent_model
+        The first part of the model taking an input and returning
+        positive activations, g(.) in the original paper.
+        Must return positive activations.
+    latent_to_logit_model
+        The second part of the model taking activation and returning
+        logits, h(.) in the original paper.
+    inputs
+        Input data of shape (n_samples, height, width, channels).
+        (x1, x2, ..., xn) in the paper.
+    labels
+        Labels of the inputs of shape (n_samples, class_id)
+    list_of_class_of_interest
+        A list of the classes id to explain. The manager will instanciate one
+        CraftTorch object per element of this list.
+    number_of_concepts
+        The number of concepts to extract. Default is 20.
+    batch_size
+        The batch size to use during training and prediction. Default is 64.
+    patch_size
+        The size of the patches (crops) to extract from the input data. Default is 64.
+    """
