Function override

[fork123aniket]

How can I override a function named explain_message() inside MessagePassing interface for each message passing layer in custom explainer which inherits from torch.nn.Module without touching each layer type' (GAT, GCN, RGCN, etc.) original implementation??

Currently, my explain_message() implements the following functionality:-

def explain_message((self, out: Tensor, x_i: Tensor, x_j: Tensor) -> Tensor:
      if out.size(-1) != self.process_message[0].normalized_shape[0]:
          out = self.convert(out)
      basis_messages = self.process_message(out)
      
      if self.message_scale is not None:
          basis_messages = basis_messages * self.message_scale.unsqueeze(
              -1)
      
          if self.message_replacement is not None:
              if basis_messages.shape == self.message_replacement.shape:
                  basis_messages = (basis_messages +
                                    (1 - self.message_scale).unsqueeze(
                                       -1) * self.message_replacement)
              else:
                  basis_messages = (basis_messages +
                                    ((1 - self.message_scale).unsqueeze(
                                        -1) *
                                     self.message_replacement.unsqueeze(
                                         0)))
      
      self.latest_messages = basis_messages
      self.latest_source_embeddings = x_j
      self.latest_target_embeddings = x_i
      
      return basis_messages

[rusty1s]

You should be able to just assign them to each layer: conv.explain_message = my_custom_explain_message.

[rusty1s]

I think I solved this in #4448. Please have a look :)

[fork123aniket]

@rusty1s This change to MessagePassing interface doesn't seem to work either. In the below custom explain_message():-

def explain_message(self, out, x_i, x_j):
    basis_messages = Sequential(LayerNorm(out.size(-1)), ReLU())(out)

    if self.message_scale is not None:
        basis_messages = basis_messages * self.message_scale.unsqueeze(
            -1)

        if self.message_replacement is not None:
            if basis_messages.shape == self.message_replacement.shape:
                basis_messages = (basis_messages +
                                  (1 - self.message_scale).unsqueeze(
                                      -1) * self.message_replacement)
            else:
                basis_messages = (basis_messages +
                                  ((1 - self.message_scale).unsqueeze(
                                      -1) *
                                   self.message_replacement.unsqueeze(
                                       0)))

    self.latest_messages = basis_messages
    self.latest_source_embeddings = x_j
    self.latest_target_embeddings = x_i

    return basis_messages

self refers to the object of GraphMaskExplainer and not of any layer's type (GAT, GCN, etc.), thereby throwing error: AttributeError: 'GCNConv' object has no attribute 'latest_source_embeddings'. I also attempted to use the test, you've written in the response to #4448 PR, but it didn't help either. Here, the problem is we need to inject message (message_scale) and learned baseline (message_replacement) in each message passing layer of the trained GNN model along with saving each of these values (basis_messages, x_j, and x_i) to their corresponding each message passing layer's object parameters (latest_messages, latest_source_embeddings, and latest_target_embeddings). I made changes to the custom __set_flags__() inside GraphMaskExplainer class, which also had gone into vein:-

def __set_flags__(self, model):
    for module in model.modules():
        if isinstance(module, MessagePassing):
            module.explain_message = self.explain_message.__get__(module, MessagePassing)
            module.explain = True
            self.message_scale = None
            self.message_replacement = None

since, here, self.message_scale only sets the GraphMaskExplainer's message_scale attribute and not sets any layer type's message_scale, which is also true for message_replacement. When I superseded self.message_scale with module.message_scale inside __set_flags__() defined above, it threw an error AttributeError: 'GCNConv' object has no attribute 'message_scale', which is also held true for message_replacement. So, this is not what we want.
The solution, at this point, I can think of is to make this custom explain_message() as MessagePassing's explain_message() default implementation since the GNNExplainer's explain_message functionality is quite easy and we can make it available within gnn_explainer.py file by overriding this function in it because GNNExplainer doesn't compute layer-wise explanation weights for each edge, however GraphMaskExplainer does. WDYT??

[rusty1s]

I am not sure I understand. self inside explain_message clearly belongs to the GNN layer if you follow the pattern in the test case. Of course, your custom explainer message should be an isolated function, not a method attached to the explainer. You also might want to check existence of attributes via if getattr(self, message_scale, None) is not None.

[fork123aniket]

Thanks, your recommendations made the code work properly. Now I'm able to produce explanations in a seamless way by overriding explain_message() of MessagePassing with the custom explain_message() defined above. To fix the problem, defined the custom explain_message() outside of the explainer class and used getattr(self, 'message_scale', None) inside this function and, bingo!!!!, it worked. Thanks so much for your quick responses and invaluable suggestions. Will soon create a PR on this!!!!

[rusty1s]

Super :) Glad that it is now working!

[688spiderbin]

@fork123aniket
Hello, thanks for your contributions,
I want to ask how to explain MessagePassing Model by using the function you changed or you can give me an example like the official colab notebook show how to explain nn.Module class model by using Captum. That will be really convenient for us to use this method to do our research.

I appreciate a lot !

[fork123aniket]

Hi @688spiderbin,

I don't know if GraphMaskExplainer has ever been made available for Captum functionality, however, you can always checkout my GitHub Repo to see an example of GraphMaskExplainer and how to explain predictions using it after April 16. Am preparing to upload the whole project as a repo, however, will be able to upload it sometime after April 16. Will notify you of this once it's available here... Thank you.