Belle II Software  release-08-01-10
GATModule Class Reference
Inheritance diagram for GATModule:
Collaboration diagram for GATModule:

Public Member Functions

def __init__ (self, in_feats, units, num_heads, in_feats_glob, use_gap=False)
 
def forward (self, graph, feat, feat_glob=None)
 

Public Attributes

 gat
 TODO.
 
 fc
 TODO.
 
 use_gap
 TODO.
 
 gap_gate
 TODO.
 
 gap
 TODO.
 

Detailed Description

Apply a GAT layer to node features, flatten outputs of attention heads
and update global features.

Definition at line 18 of file gatgap.py.

Constructor & Destructor Documentation

◆ __init__()

def __init__ (   self,
  in_feats,
  units,
  num_heads,
  in_feats_glob,
  use_gap = False 
)
Initialise the class.
:param in_feats:  TODO
:param units:  TODO
:param num_heads:  TODO
:param in_feats_glob:  TODO
:param use_gap:  TODO

Definition at line 24 of file gatgap.py.

31  ):
32  """
33  Initialise the class.
34  :param in_feats: TODO
35  :param units: TODO
36  :param num_heads: TODO
37  :param in_feats_glob: TODO
38  :param use_gap: TODO
39  """
40  super().__init__()
41 
42  self.gat = dglnn.GATConv(in_feats, units, num_heads)
43  out_feats = units * num_heads
44 
45  self.fc = torch.nn.Linear(in_feats_glob + out_feats, units)
46 
47  self.use_gap = use_gap
48  if self.use_gap:
49 
50  self.gap_gate = torch.nn.Linear(out_feats, 1)
51 
52  self.gap = GlobalAttentionPooling(self.gap_gate)
53 

Member Function Documentation

◆ forward()

def forward (   self,
  graph,
  feat,
  feat_glob = None 
)
TODO

Definition at line 54 of file gatgap.py.


The documentation for this class was generated from the following file: