Merge pull request #14 from softwaremill/fix/lrp_rules

lrp missing rules fixed

.gitignore

@@ -158,3 +158,6 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# VS Code
+.vscode/

autoxai/explainer/lrp.py

@@ -5,6 +5,7 @@
 
 import torch
 from captum.attr import LRP, LayerLRP
+from captum.attr._utils.lrp_rules import EpsilonRule, GammaRule
 
 from autoxai.explainer.base_explainer import CVExplainer
 from autoxai.explainer.model_utils import modify_modules
@@ -45,12 +46,30 @@ def calculate_features(
 
         lrp = self.create_explainer(model=model, layer=layer)
 
-        attributions = lrp.attribute(
-            input_data,
-            target=pred_label_idx,
-        )
+        attributions = lrp.attribute(input_data, target=pred_label_idx)
         return attributions
 
+    def add_rules(self, model: torch.nn.Module) -> torch.nn.Module:
+        """Add rules for the LRP explainer,
+        according to https://arxiv.org/pdf/1910.09840.pdf.
+
+        Args:
+            model: DNN object to be modified.
+
+        Returns:
+            Modified DNN object.
+        """
+        layers_number: int = len(list(model.modules()))
+        for idx_layer, module in enumerate(model.modules()):
+            if idx_layer <= layers_number // 2:
+                setattr(module, "rule", GammaRule())
+            elif idx_layer != (layers_number - 1):
+                setattr(module, "rule", EpsilonRule())
+            else:
+                setattr(module, "rule", EpsilonRule(epsilon=0))  # LRP-0
+
+        return model
+
 
 class LRPCVExplainer(BaseLRPCVExplainer):
     """LRP algorithm explainer."""
@@ -68,7 +87,7 @@ def create_explainer(self, **kwargs) -> Union[LRP, LayerLRP]:
         if model is None:
             raise RuntimeError(f"Missing or `None` argument `model` passed: {kwargs}")
 
-        model = modify_modules(model)
+        model = self.add_rules(modify_modules(model))
 
         return LRP(model=model)
 
@@ -92,6 +111,6 @@ def create_explainer(self, **kwargs) -> Union[LRP, LayerLRP]:
                 f"Missing or `None` arguments `model` or `layer` passed: {kwargs}"
             )
 
-        model = modify_modules(model)
+        model = self.add_rules(modify_modules(model))
 
         return LayerLRP(model=model, layer=layer)

autoxai/explainer/model_utils.py

@@ -1,5 +1,4 @@
 """File contains functions to modifiy DNN models."""
-
 import torch