Points wise importance/attention during points cloud classification

[vctorwei]

Dear pyg-team

Task
I have a points cloud dataset of tumours, the dataset is used for the points cloud level classification
I wish to find boundary areas (points with high attention/importance ) that are important in the classification task.

I learnt that GNNExplainer is capable of this task, but I wish to obtain the important points during the training so that I can use them to train another model simultaneously.

Possible Solution
I also learnt GlobalAttention is a pooling method that can export node importance in graph classification task, but I think this method only applies to the last layer downsampled points cloud graph ? (e.g. 1024 points --> edge_conv + fps --> 128 points --> edge_conv + fps --> 32 points --> global attention --> 16 of 32 points have importance larger than ...) .

How can I obtain the importance of all 1024 points in the original points cloud? Is there a way to reverse the fps and assign those 32 points importance to their neighbours?

Another question is can I just extract the importance by returning the softmax output gate here?

        gate = softmax(gate, batch, num_nodes=size)
        out = scatter_add(gate * x, batch, dim=0, dim_size=size)
        return out, gate

Any other possible solutions?

Thank you
Wei

[rusty1s]

If you want to learn the importance of nodes jointly with the prediction task, you could apply a soft masking approach at the beginning of your model:

self.mask = torch.nn.Parameter(torch.randn(num_nodes, 1))

def forward(self, x, edge_index, ...):
    x = self.mask.sigmoid() * x
    ...

Let me know if this is helpful to you :)
...

[vctorwei]

Thank you!
Does the following code look correct? num_node is depending on the batch size so I'm not sure

    def __set_masks__(self, x):
        num_node=x.size()[0]
        self.node_mask = torch.nn.Parameter(torch.randn(num_node, 1)).to(device)
    def forward(self, x, edge_index, edge_weight, pos, batch):
        self.__set_masks__(x)
        node_importance = self.node_mask.sigmoid()
        x = node_importance * x
        ...
       return x, node_importance

[rusty1s]

Mh, I see. If you have multiple point clouds, this will not work, as you correctly pointed out, will need to learn node-dependent masks.

[vctorwei]

Can you please give me an example of how to do that in batch with whole points cloud/graph level classification task?

[rusty1s]

I don't think this is possible, at least not in the early stages of a network (since you do not have meaningful features to filter out nodes early on). As such, using GlobalAttention sounds indeed like your best option.

[vctorwei]

Thank you for the reply! :)
I guess I'll just use the GlobalAttention and then assign the attention of the sampled points to their neighbours in the original points cloud using nearest