[bug][doc] Fix bug in batched heterograph training documentation (#2204)

* [bug][doc] Fix bug in batched heterograph training documentation

* added references of "above" with hyperlinks

* address comments

* fix API doc on EdgeDataLoader
Co-authored-by: Mufei Li <mufeili1996@gmail.com>

docs/source/guide/training-edge.rst

@@ -51,7 +51,8 @@ node representations on each edge.
     import dgl.function as fn
     class DotProductPredictor(nn.Module):
         def forward(self, graph, h):
-            # h contains the node representations computed from the GNN above.
+            # h contains the node representations computed from the GNN defined
+            # in the node classification section (Section 5.1).
             with graph.local_scope():
                 graph.ndata['h'] = h
                 graph.apply_edges(fn.u_dot_v('h', 'h', 'score'))
@@ -75,7 +76,8 @@ e.g. as logits of a categorical distribution.
             return {'score': score}
     
         def forward(self, graph, h):
-            # h contains the node representations computed from the GNN above.
+            # h contains the node representations computed from the GNN defined
+            # in the node classification section (Section 5.1).
             with graph.local_scope():
                 graph.ndata['h'] = h
                 graph.apply_edges(self.apply_edges)
@@ -122,6 +124,7 @@ not include early stopping and model saving.
         opt.step()
         print(loss.item())
 
+.. _guide-training-edge-classification-heterogeneous-graph:
 
 Heterogeneous graph
 ~~~~~~~~~~~~~~~~~~~
@@ -141,7 +144,8 @@ heterogeneous graph, you only need to specify the edge type in
     class HeteroDotProductPredictor(nn.Module):
         def forward(self, graph, h, etype):
             # h contains the node representations for each edge type computed from
-            # the GNN above.
+            # the GNN for heterogeneous graphs defined in the node classification
+            # section (Section 5.1).
             with graph.local_scope():
                 graph.ndata['h'] = h   # assigns 'h' of all node types in one shot
                 graph.apply_edges(fn.u_dot_v('h', 'h', 'score'), etype=etype)
@@ -163,7 +167,9 @@ You can similarly write a ``HeteroMLPPredictor``.
             return {'score': score}
     
         def forward(self, graph, h, etype):
-            # h contains the node representations computed from the GNN above.
+            # h contains the node representations for each edge type computed from
+            # the GNN for heterogeneous graphs defined in the node classification
+            # section (Section 5.1).
             with graph.local_scope():
                 graph.ndata['h'] = h   # assigns 'h' of all node types in one shot
                 graph.apply_edges(self.apply_edges, etype=etype)
@@ -217,16 +223,18 @@ Predicting Edge Type of an Existing Edge on a Heterogeneous Graph
 Sometimes you may want to predict which type an existing edge belongs
 to.
 
-For instance, given the heterogeneous graph above, your task is given an
-edge connecting a user and an item, predict whether the user would
-``click`` or ``dislike`` an item.
+For instance, given the
+:ref:`heterogeneous graph example <guide-training-heterogeneous-graph-example>`,
+your task is given an edge connecting a user and an item, to predict whether
+the user would ``click`` or ``dislike`` an item.
 
 This is a simplified version of rating prediction, which is common in
 recommendation literature.
 
 You can use a heterogeneous graph convolution network to obtain the node
-representations. For instance, you can still use the RGCN above for this
-purpose.
+representations. For instance, you can still use the
+:ref:`RGCN defined previously <guide-training-rgcn-node-classification>`
+for this purpose.
 
 To predict the type of an edge, you can simply repurpose the
 ``HeteroDotProductPredictor`` above so that it takes in another graph
@@ -235,14 +243,17 @@ predicted, and emits the score of each type for every edge.
 
 In the example here, you will need a graph that has two node types
 ``user`` and ``item``, and one single edge type that “merges” all the
-edge types from ``user`` and ``item``, i.e. ``like`` and ``dislike``.
-This can be conveniently created using
-:meth:`relation slicing <dgl.DGLHeteroGraph.__getitem__>`.
+edge types from ``user`` and ``item``, i.e. ``click`` and ``dislike``.
+This can be conveniently created using the following syntax:
 
 .. code:: python
 
     dec_graph = hetero_graph['user', :, 'item']
 
+which returns a heterogeneous graphs with node type ``user`` and ``item``,
+as well as a single edge type combining all edge types in between, i.e.
+``click`` and ``dislike``.
+
 Since the statement above also returns the original edge types as a
 feature named ``dgl.ETYPE``, we can use that as labels.
 
@@ -267,7 +278,8 @@ can write your predictor module as follows.
     
         def forward(self, graph, h):
             # h contains the node representations for each edge type computed from
-            # the GNN above.
+            # the GNN for heterogeneous graphs defined in the node classification
+            # section (Section 5.1).
             with graph.local_scope():
                 graph.ndata['h'] = h   # assigns 'h' of all node types in one shot
                 graph.apply_edges(self.apply_edges)

docs/source/guide/training-graph.rst

@@ -233,14 +233,14 @@ representations for each node type.
     class RGCN(nn.Module):
         def __init__(self, in_feats, hid_feats, out_feats, rel_names):
             super().__init__()
-            
+    
             self.conv1 = dglnn.HeteroGraphConv({
                 rel: dglnn.GraphConv(in_feats, hid_feats)
                 for rel in rel_names}, aggregate='sum')
             self.conv2 = dglnn.HeteroGraphConv({
                 rel: dglnn.GraphConv(hid_feats, out_feats)
                 for rel in rel_names}, aggregate='sum')
-      
+    
         def forward(self, graph, inputs):
             # inputs are features of nodes
             h = self.conv1(graph, inputs)
@@ -251,21 +251,13 @@ representations for each node type.
     class HeteroClassifier(nn.Module):
         def __init__(self, in_dim, hidden_dim, n_classes, rel_names):
             super().__init__()
-            
-            self.conv1 = dglnn.HeteroGraphConv({
-                rel: dglnn.GraphConv(in_feats, hid_feats)
-                for rel in rel_names}, aggregate='sum')
-            self.conv2 = dglnn.HeteroGraphConv({
-                rel: dglnn.GraphConv(hid_feats, out_feats)
-                for rel in rel_names}, aggregate='sum')
+
+            self.rgcn = RGCN(in_dim, hidden_dim, hidden_dim, rel_names)
             self.classify = nn.Linear(hidden_dim, n_classes)
     
         def forward(self, g):
             h = g.ndata['feat']
-            # Apply graph convolution and activation.
-            h = F.relu(self.conv1(g, h))
-            h = F.relu(self.conv2(g, h))
-    
+            h = self.rgcn(g, h)
             with g.local_scope():
                 g.ndata['h'] = h
                 # Calculate graph representation by average readout.
@@ -278,7 +270,8 @@ The rest of the code is not different from that for homogeneous graphs.
 
 .. code:: python
 
-    model = HeteroClassifier(10, 20, 5)
+    # etypes is the list of edge types as strings.
+    model = HeteroClassifier(10, 20, 5, etypes)
     opt = torch.optim.Adam(model.parameters())
     for epoch in range(20):
         for batched_graph, labels in dataloader:

docs/source/guide/training-link.rst

@@ -52,7 +52,8 @@ Model Implementation Difference from Edge Classification
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 The neural network model to compute the score between :math:`u` and
-:math:`v` is identical to the edge regression model described above.
+:math:`v` is identical to the edge regression model described
+:ref:`above <guide-training-edge-classification>`.
 
 Here is an example of using dot product to compute the scores on edges.
 
@@ -60,7 +61,8 @@ Here is an example of using dot product to compute the scores on edges.
 
     class DotProductPredictor(nn.Module):
         def forward(self, graph, h):
-            # h contains the node representations computed from the GNN above.
+            # h contains the node representations computed from the GNN defined
+            # in the node classification section (Section 5.1).
             with graph.local_scope():
                 graph.ndata['h'] = h
                 graph.apply_edges(fn.u_dot_v('h', 'h', 'score'))
@@ -143,15 +145,16 @@ Link prediction on heterogeneous graphs is not very different from that
 on homogeneous graphs. The following assumes that we are predicting on
 one edge type, and it is easy to extend it to multiple edge types.
 
-For example, you can reuse the ``HeteroDotProductPredictor`` above for
-computing the scores of the edges of an edge type for link prediction.
+For example, you can reuse the ``HeteroDotProductPredictor``
+:ref:`above <guide-training-edge-classification-heterogeneous-graph>`
+for computing the scores of the edges of an edge type for link prediction.
 
 .. code:: python
 
     class HeteroDotProductPredictor(nn.Module):
         def forward(self, graph, h, etype):
-            # h contains the node representations for each edge type computed from
-            # the GNN above.
+            # h contains the node representations for each node type computed from
+            # the GNN defined in the previous section (Section 5.1).
             with graph.local_scope():
                 graph.ndata['h'] = h
                 graph.apply_edges(fn.u_dot_v('h', 'h', 'score'), etype=etype)

docs/source/guide/training-node.rst

@@ -179,7 +179,8 @@ all node types.
 feature tensors as input, and returns another dictionary of node types
 and node features.
 
-So given that we have the user and item features in the example above.
+So given that we have the user and item features in the
+:ref:`heterogeneous graph example <guide-training-heterogeneous-graph-example>`.
 
 .. code:: python
 

docs/source/guide/training.rst

@@ -31,6 +31,8 @@ get the graph for a single-graph dataset with something like
 
 Note: In this chapter we will use PyTorch as backend.
 
+.. _guide-training-heterogeneous-graph-example:
+
 Heterogeneous Graphs
 ~~~~~~~~~~~~~~~~~~~~