[Feature] Add mhspmm to speed up GAT (#234)

* new multihead kernel * new gat layer * update * mhspmm complete * Modify GATLayer and fix bugs * fix bugs * fix bugs in tests Co-authored-by: fishmingyu <fishmingyu@github.com> Co-authored-by: THINK2TRY <48679723+THINK2TRY@users.noreply.github.com>
THUDM · May 27, 2021 · 9f2404f · 9f2404f
1 parent bb25aa4
commit 9f2404f
Show file tree

Hide file tree

Showing 16 changed files with 1,855 additions and 1,425 deletions.
diff --git a/cogdl/models/agc/daegc.py b/cogdl/models/agc/daegc.py
@@ -49,12 +49,8 @@ def __init__(self, num_features, hidden_size, embedding_size, num_heads, dropout
         self.embedding_size = embedding_size
         self.dropout = dropout
         self.num_clusters = num_clusters
-        self.att1 = GATLayer(
-            num_features, hidden_size, dropout=dropout, alpha=0.2, nhead=num_heads, concat=True, fast_mode=False
-        )
-        self.att2 = GATLayer(
-            hidden_size * num_heads, embedding_size, dropout=dropout, alpha=0.2, nhead=1, concat=False, fast_mode=False
-        )
+        self.att1 = GATLayer(num_features, hidden_size, dropout=dropout, alpha=0.2, nhead=num_heads, concat=True)
+        self.att2 = GATLayer(hidden_size * num_heads, embedding_size, dropout=dropout, alpha=0.2, nhead=1, concat=False)
         self.cluster_center = None
 
     def get_trainer(self, task, args):