[Misc] Black auto fix. (dmlc#4679)

Co-authored-by: Steve <[email protected]>

Author: frozenbugs

examples/pytorch/TAHIN/TAHIN.py

@@ -6,187 +6,241 @@
 import dgl.function as fn
 from dgl.nn.pytorch import GATConv
 
-#Semantic attention in the metapath-based aggregation (the same as that in the HAN)
+
+# Semantic attention in the metapath-based aggregation (the same as that in the HAN)
 class SemanticAttention(nn.Module):
     def __init__(self, in_size, hidden_size=128):
         super(SemanticAttention, self).__init__()
 
         self.project = nn.Sequential(
             nn.Linear(in_size, hidden_size),
             nn.Tanh(),
-            nn.Linear(hidden_size, 1, bias=False)
+            nn.Linear(hidden_size, 1, bias=False),
         )
 
     def forward(self, z):
-        '''
+        """
         Shape of z: (N, M , D*K)
         N: number of nodes
         M: number of metapath patterns
         D: hidden_size
         K: number of heads
-        '''
-        w = self.project(z).mean(0)                    # (M, 1)
-        beta = torch.softmax(w, dim=0)                 # (M, 1)
-        beta = beta.expand((z.shape[0],) + beta.shape) # (N, M, 1)
+        """
+        w = self.project(z).mean(0)  # (M, 1)
+        beta = torch.softmax(w, dim=0)  # (M, 1)
+        beta = beta.expand((z.shape[0],) + beta.shape)  # (N, M, 1)
+
+        return (beta * z).sum(1)  # (N, D * K)
 
-        return (beta * z).sum(1)                       # (N, D * K)
 
-#Metapath-based aggregation (the same as the HANLayer)
+# Metapath-based aggregation (the same as the HANLayer)
 class HANLayer(nn.Module):
-    def __init__(self, meta_path_patterns, in_size, out_size, layer_num_heads, dropout):
+    def __init__(
+        self, meta_path_patterns, in_size, out_size, layer_num_heads, dropout
+    ):
         super(HANLayer, self).__init__()
 
         # One GAT layer for each meta path based adjacency matrix
         self.gat_layers = nn.ModuleList()
         for i in range(len(meta_path_patterns)):
-            self.gat_layers.append(GATConv(in_size, out_size, layer_num_heads,
-                                           dropout, dropout, activation=F.elu,
-                                           allow_zero_in_degree=True))
-        self.semantic_attention = SemanticAttention(in_size=out_size * layer_num_heads)
-        self.meta_path_patterns = list(tuple(meta_path_pattern) for meta_path_pattern in meta_path_patterns)
+            self.gat_layers.append(
+                GATConv(
+                    in_size,
+                    out_size,
+                    layer_num_heads,
+                    dropout,
+                    dropout,
+                    activation=F.elu,
+                    allow_zero_in_degree=True,
+                )
+            )
+        self.semantic_attention = SemanticAttention(
+            in_size=out_size * layer_num_heads
+        )
+        self.meta_path_patterns = list(
+            tuple(meta_path_pattern) for meta_path_pattern in meta_path_patterns
+        )
 
         self._cached_graph = None
         self._cached_coalesced_graph = {}
 
     def forward(self, g, h):
         semantic_embeddings = []
-        #obtain metapath reachable graph
+        # obtain metapath reachable graph
         if self._cached_graph is None or self._cached_graph is not g:
             self._cached_graph = g
             self._cached_coalesced_graph.clear()
             for meta_path_pattern in self.meta_path_patterns:
-                self._cached_coalesced_graph[meta_path_pattern] = dgl.metapath_reachable_graph(
-                        g, meta_path_pattern)
+                self._cached_coalesced_graph[
+                    meta_path_pattern
+                ] = dgl.metapath_reachable_graph(g, meta_path_pattern)
 
         for i, meta_path_pattern in enumerate(self.meta_path_patterns):
             new_g = self._cached_coalesced_graph[meta_path_pattern]
             semantic_embeddings.append(self.gat_layers[i](new_g, h).flatten(1))
-        semantic_embeddings = torch.stack(semantic_embeddings, dim=1)                  # (N, M, D * K)
+        semantic_embeddings = torch.stack(
+            semantic_embeddings, dim=1
+        )  # (N, M, D * K)
+
+        return self.semantic_attention(semantic_embeddings)  # (N, D * K)
 
-        return self.semantic_attention(semantic_embeddings)                            # (N, D * K)
 
-#Relational neighbor aggregation
+# Relational neighbor aggregation
 class RelationalAGG(nn.Module):
     def __init__(self, g, in_size, out_size, dropout=0.1):
         super(RelationalAGG, self).__init__()
         self.in_size = in_size
         self.out_size = out_size
 
-        #Transform weights for different types of edges
-        self.W_T = nn.ModuleDict({
-            name : nn.Linear(in_size, out_size, bias = False) for name in g.etypes
-        })
+        # Transform weights for different types of edges
+        self.W_T = nn.ModuleDict(
+            {
+                name: nn.Linear(in_size, out_size, bias=False)
+                for name in g.etypes
+            }
+        )
 
-        #Attention weights for different types of edges
-        self.W_A = nn.ModuleDict({
-            name : nn.Linear(out_size, 1, bias = False) for name in g.etypes
-        })
+        # Attention weights for different types of edges
+        self.W_A = nn.ModuleDict(
+            {name: nn.Linear(out_size, 1, bias=False) for name in g.etypes}
+        )
 
-        #layernorm
+        # layernorm
         self.layernorm = nn.LayerNorm(out_size)
 
-        #dropout layer
+        # dropout layer
         self.dropout = nn.Dropout(dropout)
 
     def forward(self, g, feat_dict):
-        funcs={}
+        funcs = {}
         for srctype, etype, dsttype in g.canonical_etypes:
-            g.nodes[dsttype].data['h'] = feat_dict[dsttype] #nodes' original feature
-            g.nodes[srctype].data['h'] = feat_dict[srctype]  
-            g.nodes[srctype].data['t_h'] = self.W_T[etype](feat_dict[srctype])  #src nodes' transformed feature
-
-            #compute the attention numerator (exp)
-            g.apply_edges(fn.u_mul_v('t_h','h','x'),etype=etype)
-            g.edges[etype].data['x'] = torch.exp(self.W_A[etype](g.edges[etype].data['x']))
-
-            #first update to compute the attention denominator (\sum exp)
-            funcs[etype] = (fn.copy_e('x', 'm'), fn.sum('m', 'att'))  
-        g.multi_update_all(funcs, 'sum') 
-
-        funcs={}
+            g.nodes[dsttype].data["h"] = feat_dict[
+                dsttype
+            ]  # nodes' original feature
+            g.nodes[srctype].data["h"] = feat_dict[srctype]
+            g.nodes[srctype].data["t_h"] = self.W_T[etype](
+                feat_dict[srctype]
+            )  # src nodes' transformed feature
+
+            # compute the attention numerator (exp)
+            g.apply_edges(fn.u_mul_v("t_h", "h", "x"), etype=etype)
+            g.edges[etype].data["x"] = torch.exp(
+                self.W_A[etype](g.edges[etype].data["x"])
+            )
+
+            # first update to compute the attention denominator (\sum exp)
+            funcs[etype] = (fn.copy_e("x", "m"), fn.sum("m", "att"))
+        g.multi_update_all(funcs, "sum")
+
+        funcs = {}
         for srctype, etype, dsttype in g.canonical_etypes:
-            g.apply_edges(fn.e_div_v('x', 'att', 'att'),etype=etype) #compute attention weights (numerator/denominator)
-            funcs[etype] = (fn.u_mul_e('h', 'att', 'm'), fn.sum('m', 'h'))  #\sum(h0*att) -> h1
-        #second update to obtain h1
-        g.multi_update_all(funcs, 'sum') 
-
-        #apply activation, layernorm, and dropout
-        feat_dict={}
+            g.apply_edges(
+                fn.e_div_v("x", "att", "att"), etype=etype
+            )  # compute attention weights (numerator/denominator)
+            funcs[etype] = (
+                fn.u_mul_e("h", "att", "m"),
+                fn.sum("m", "h"),
+            )  # \sum(h0*att) -> h1
+        # second update to obtain h1
+        g.multi_update_all(funcs, "sum")
+
+        # apply activation, layernorm, and dropout
+        feat_dict = {}
         for ntype in g.ntypes:
-            feat_dict[ntype] = self.dropout(self.layernorm(F.relu_(g.nodes[ntype].data['h']))) #apply activation, layernorm, and dropout
-        
+            feat_dict[ntype] = self.dropout(
+                self.layernorm(F.relu_(g.nodes[ntype].data["h"]))
+            )  # apply activation, layernorm, and dropout
+
         return feat_dict
 
+
 class TAHIN(nn.Module):
-    def __init__(self, g, meta_path_patterns, in_size, out_size, num_heads, dropout):
+    def __init__(
+        self, g, meta_path_patterns, in_size, out_size, num_heads, dropout
+    ):
         super(TAHIN, self).__init__()
 
-        #embeddings for different types of nodes, h0
+        # embeddings for different types of nodes, h0
         self.initializer = nn.init.xavier_uniform_
-        self.feature_dict = nn.ParameterDict({
-            ntype: nn.Parameter(self.initializer(torch.empty(g.num_nodes(ntype), in_size))) for ntype in g.ntypes
-        })
+        self.feature_dict = nn.ParameterDict(
+            {
+                ntype: nn.Parameter(
+                    self.initializer(torch.empty(g.num_nodes(ntype), in_size))
+                )
+                for ntype in g.ntypes
+            }
+        )
 
-        #relational neighbor aggregation, this produces h1
+        # relational neighbor aggregation, this produces h1
         self.RelationalAGG = RelationalAGG(g, in_size, out_size)
 
-        #metapath-based aggregation modules for user and item, this produces h2
-        self.meta_path_patterns = meta_path_patterns 
-        #one HANLayer for user, one HANLayer for item
-        self.hans = nn.ModuleDict({
-            key: HANLayer(value, in_size, out_size, num_heads, dropout) for key, value in self.meta_path_patterns.items()
-        })
-
-        #layers to combine h0, h1, and h2
-        #used to update node embeddings
-        self.user_layer1 = nn.Linear((num_heads+1)*out_size, out_size, bias=True)
-        self.user_layer2 = nn.Linear(2*out_size, out_size, bias=True)
-        self.item_layer1 = nn.Linear((num_heads+1)*out_size, out_size, bias=True)
-        self.item_layer2 = nn.Linear(2*out_size, out_size, bias=True)
-
-        #layernorm
+        # metapath-based aggregation modules for user and item, this produces h2
+        self.meta_path_patterns = meta_path_patterns
+        # one HANLayer for user, one HANLayer for item
+        self.hans = nn.ModuleDict(
+            {
+                key: HANLayer(value, in_size, out_size, num_heads, dropout)
+                for key, value in self.meta_path_patterns.items()
+            }
+        )
+
+        # layers to combine h0, h1, and h2
+        # used to update node embeddings
+        self.user_layer1 = nn.Linear(
+            (num_heads + 1) * out_size, out_size, bias=True
+        )
+        self.user_layer2 = nn.Linear(2 * out_size, out_size, bias=True)
+        self.item_layer1 = nn.Linear(
+            (num_heads + 1) * out_size, out_size, bias=True
+        )
+        self.item_layer2 = nn.Linear(2 * out_size, out_size, bias=True)
+
+        # layernorm
         self.layernorm = nn.LayerNorm(out_size)
 
-        #network to score the node pairs
+        # network to score the node pairs
         self.pred = nn.Linear(out_size, out_size)
         self.dropout = nn.Dropout(dropout)
         self.fc = nn.Linear(out_size, 1)
 
     def forward(self, g, user_key, item_key, user_idx, item_idx):
-        #relational neighbor aggregation, h1
+        # relational neighbor aggregation, h1
         h1 = self.RelationalAGG(g, self.feature_dict)
 
-        #metapath-based aggregation, h2
+        # metapath-based aggregation, h2
         h2 = {}
         for key in self.meta_path_patterns.keys():
             h2[key] = self.hans[key](g, self.feature_dict[key])
 
-        #update node embeddings
+        # update node embeddings
         user_emb = torch.cat((h1[user_key], h2[user_key]), 1)
         item_emb = torch.cat((h1[item_key], h2[item_key]), 1)
         user_emb = self.user_layer1(user_emb)
         item_emb = self.item_layer1(item_emb)
-        user_emb = self.user_layer2(torch.cat((user_emb, self.feature_dict[user_key]), 1))
-        item_emb = self.item_layer2(torch.cat((item_emb, self.feature_dict[item_key]), 1))
+        user_emb = self.user_layer2(
+            torch.cat((user_emb, self.feature_dict[user_key]), 1)
+        )
+        item_emb = self.item_layer2(
+            torch.cat((item_emb, self.feature_dict[item_key]), 1)
+        )
 
-        #Relu
+        # Relu
         user_emb = F.relu_(user_emb)
         item_emb = F.relu_(item_emb)
-        
-        #layer norm
+
+        # layer norm
         user_emb = self.layernorm(user_emb)
         item_emb = self.layernorm(item_emb)
-        
-        #obtain users/items embeddings and their interactions
+
+        # obtain users/items embeddings and their interactions
         user_feat = user_emb[user_idx]
         item_feat = item_emb[item_idx]
-        interaction = user_feat*item_feat
+        interaction = user_feat * item_feat
 
-        #score the node pairs
+        # score the node pairs
         pred = self.pred(interaction)
-        pred = self.dropout(pred) #dropout
+        pred = self.dropout(pred)  # dropout
         pred = self.fc(pred)
         pred = torch.sigmoid(pred)
 
         return pred.squeeze(1)
-