get_attention parameter in GlobalAttentionPooling

python/dgl/nn/pytorch/glob.py

@@ -419,7 +419,7 @@ def __init__(self, gate_nn, feat_nn=None):
         self.gate_nn = gate_nn
         self.feat_nn = feat_nn
 
-    def forward(self, graph, feat):
+    def forward(self, graph, feat, get_attention=False):
         r"""
 
         Compute global attention pooling.
@@ -431,12 +431,17 @@ def forward(self, graph, feat):
         feat : torch.Tensor
             The input node feature with shape :math:`(N, D)` where :math:`N` is the
             number of nodes in the graph, and :math:`D` means the size of features.
+        get_attention : bool, optional
+            Whether to return the attention values from gate_nn. Default to False.
 
         Returns
         -------
         torch.Tensor
             The output feature with shape :math:`(B, D)`, where :math:`B` refers
             to the batch size.
+        torch.Tensor, optional
+            The attention values of shape :math:`(N, 1)`, where :math:`N` is the number of
+            nodes in the graph. This is returned only when :attr:`get_attention` is ``True``.
         """
         with graph.local_scope():
             gate = self.gate_nn(feat)
@@ -450,8 +455,11 @@ def forward(self, graph, feat):
             graph.ndata['r'] = feat * gate
             readout = sum_nodes(graph, 'r')
             graph.ndata.pop('r')
-
-            return readout
+
+            if get_attention:
+                return readout, gate
+            else:
+                return readout
 
 
 class Set2Set(nn.Module):