[Example] layers number fix in GAT example (#3872)

* gat layers fix

* minor bug
Co-authored-by: decoherencer <decoherencer@users.noreply.github.com>
Co-authored-by: Minjie Wang <wmjlyjemaine@gmail.com>

examples/pytorch/gat/gat.py

@@ -31,20 +31,25 @@ class GAT(nn.Module):
         self.num_layers = num_layers
         self.gat_layers = nn.ModuleList()
         self.activation = activation
+        if num_layers > 1:
         # input projection (no residual)
-        self.gat_layers.append(GATConv(
-            in_dim, num_hidden, heads[0],
-            feat_drop, attn_drop, negative_slope, False, self.activation))
-        # hidden layers
-        for l in range(1, num_layers):
-            # due to multi-head, the in_dim = num_hidden * num_heads
             self.gat_layers.append(GATConv(
-                num_hidden * heads[l-1], num_hidden, heads[l],
-                feat_drop, attn_drop, negative_slope, residual, self.activation))
-        # output projection
-        self.gat_layers.append(GATConv(
-            num_hidden * heads[-2], num_classes, heads[-1],
-            feat_drop, attn_drop, negative_slope, residual, None))
+                in_dim, num_hidden, heads[0],
+                feat_drop, attn_drop, negative_slope, False, self.activation))
+            # hidden layers
+            for l in range(1, num_layers-1):
+                # due to multi-head, the in_dim = num_hidden * num_heads
+                self.gat_layers.append(GATConv(
+                    num_hidden * heads[l-1], num_hidden, heads[l],
+                    feat_drop, attn_drop, negative_slope, residual, self.activation))
+            # output projection
+            self.gat_layers.append(GATConv(
+                num_hidden * heads[-2], num_classes, heads[-1],
+                feat_drop, attn_drop, negative_slope, residual, None))
+        else:
+            self.gat_layers.append(GATConv(
+                in_dim, num_classes, heads[0],
+                feat_drop, attn_drop, negative_slope, residual, None))
 
     def forward(self, inputs):
         h = inputs

examples/pytorch/gat/train.py

@@ -79,7 +79,7 @@ def main(args):
     g = dgl.add_self_loop(g)
     n_edges = g.number_of_edges()
     # create model
-    heads = ([args.num_heads] * args.num_layers) + [args.num_out_heads]
+    heads = ([args.num_heads] * (args.num_layers-1)) + [args.num_out_heads]
     model = GAT(g,
                 args.num_layers,
                 num_feats,
@@ -153,7 +153,7 @@ if __name__ == '__main__':
                         help="number of hidden attention heads")
     parser.add_argument("--num-out-heads", type=int, default=1,
                         help="number of output attention heads")
-    parser.add_argument("--num-layers", type=int, default=1,
+    parser.add_argument("--num-layers", type=int, default=2,
                         help="number of hidden layers")
     parser.add_argument("--num-hidden", type=int, default=8,
                         help="number of hidden units")

examples/pytorch/gat/train_ppi.py

@@ -57,7 +57,7 @@ def main(args):
     n_classes = train_dataset.num_labels
     num_feats = g.ndata['feat'].shape[1]
     g = g.int().to(device)
-    heads = ([args.num_heads] * args.num_layers) + [args.num_out_heads]
+    heads = ([args.num_heads] * (args.num_layers-1)) + [args.num_out_heads]
     # define the model
     model = GAT(g,
                 args.num_layers,
@@ -129,7 +129,7 @@ if __name__ == '__main__':
                         help="number of hidden attention heads")
     parser.add_argument("--num-out-heads", type=int, default=6,
                         help="number of output attention heads")
-    parser.add_argument("--num-layers", type=int, default=2,
+    parser.add_argument("--num-layers", type=int, default=3,
                         help="number of hidden layers")
     parser.add_argument("--num-hidden", type=int, default=256,
                         help="number of hidden units")