fix the bug causes avs_projected grad lost

Author: yztxwd

test/test_cav_trainer.py

@@ -1,13 +1,15 @@
 import unittest
+from functools import partial
 from pathlib import Path
 
 import torch
 from Bio import motifs as Bio_motifs
+from captum.attr import DeepLift
 
 from tpcav import helper
 from tpcav.cavs import CavTrainer
 from tpcav.concepts import ConceptBuilder
-from tpcav.tpcav_model import TPCAV
+from tpcav.tpcav_model import TPCAV, _abs_attribution_func
 
 
 class DummyModelSeq(torch.nn.Module):
@@ -22,6 +24,11 @@ def forward(self, seq):
         y_hat = self.layer2(y_hat)
         return y_hat
 
+    def foward_from_layer1(self, y_hat):
+        y_hat = y_hat.squeeze(-1)
+        y_hat = self.layer2(y_hat)
+        return y_hat
+
 
 class DummyModelSeqChrom(torch.nn.Module):
     def __init__(self):
@@ -174,6 +181,49 @@ def pack_data_iters(df):
             ],
         )
 
+        # compute layer attributions using the old way
+        random1_avs = []
+        random2_avs = []
+        for inputs in pack_data_iters(random_regions_1):
+            av = tpcav_model._layer_output(*[i.to(tpcav_model.device) for i in inputs])
+            random1_avs.append(av.detach().cpu())
+        for inputs in pack_data_iters(random_regions_2):
+            av = tpcav_model._layer_output(*[i.to(tpcav_model.device) for i in inputs])
+            random2_avs.append(av.detach().cpu())
+        random1_avs = torch.cat(random1_avs, dim=0)
+        random2_avs = torch.cat(random2_avs, dim=0)
+
+        random1_avs_residual, random1_avs_projected = tpcav_model.project_activations(
+            random1_avs
+        )
+        random2_avs_residual, random2_avs_projected = tpcav_model.project_activations(
+            random2_avs
+        )
+
+        def forward_from_layer_1_embeddings(tm, avs_residual, avs_projected):
+            y_hat = tm.embedding_to_layer_activation(avs_residual, avs_projected)
+            y_hat = tm.model.foward_from_layer1(y_hat)
+            return y_hat
+
+        tpcav_model.forward = partial(forward_from_layer_1_embeddings, tpcav_model)
+
+        dl = DeepLift(tpcav_model)
+        attributions_old = dl.attribute(
+            (
+                random1_avs_residual.to(tpcav_model.device),
+                random1_avs_projected.to(tpcav_model.device),
+            ),
+            baselines=(
+                random2_avs_residual.to(tpcav_model.device),
+                random2_avs_projected.to(tpcav_model.device),
+            ),
+            custom_attribution_func=_abs_attribution_func,
+        )
+        attr_residual, attr_projected = attributions_old
+        attributions_old = torch.cat((attr_projected, attr_residual), dim=1)
+
+        self.assertTrue(torch.allclose(attributions.cpu(), attributions_old.cpu()))
+
 
 if __name__ == "__main__":
     unittest.main()

tpcav/tpcav_model.py

@@ -1,4 +1,3 @@
-import inspect
 import logging
 from functools import partial
 from typing import Dict, Iterable, List, Optional, Tuple
@@ -12,6 +11,10 @@
 
 def _abs_attribution_func(multipliers, inputs, baselines):
     "Multiplier x abs(inputs - baselines) to avoid double-sign effects."
+    # print(f"inputs: {inputs[1][:5]}")
+    # print(f"baselines: {baselines[1][:5]}")
+    # print(f"multipliers: {multipliers[0][:5]}")
+    # print(f"multipliers: {multipliers[1][:5]}")
     return tuple(
         (input_ - baseline).abs() * multiplier
         for input_, baseline, multiplier in zip(inputs, baselines, multipliers)
@@ -61,7 +64,6 @@ def restore_tpcav_state(self, tpcav_state_dict: Dict) -> None:
         self._set_buffer("pca_inv", tpcav_state_dict["pca_inv"])
         self._set_buffer("orig_shape", tpcav_state_dict["orig_shape"])
         self.fitted = True
-        print(inspect.currentframe().f_back.f_code.co_name)
         logger.warning(
             "Restored TPCAV state, please set model attribute!\n\n Example: self.model = Model_class()",
         )
@@ -190,7 +192,11 @@ def layer_attributions(
             bavs = self._layer_output(*[bi.to(self.device) for bi in binputs])
             bavs_residual, bavs_projected = self.project_activations(bavs)
 
+            # detach the projected tensor as it's connnected to the original input graph,
+            # detaching it would keep the gradients on it
             if avs_projected is not None:
+                avs_projected = avs_projected.detach()
+                bavs_projected = bavs_projected.detach()
                 attribution = deeplift.attribute(
                     (avs_residual.to(self.device), avs_projected.to(self.device)),
                     baselines=(
@@ -258,7 +264,6 @@ def input_attributions(
 
         attributions = []
         for inputs, binputs in zip(target_batches, baseline_batches):
-
             attribution = deeplift.attribute(
                 tuple([i.to(self.device) for i in inputs]),
                 baselines=tuple([bi.to(self.device) for bi in binputs]),