📝 Add PyDocstrings for GCN Tutorials model.py and train.py

docs/source/index.rst

@@ -32,7 +32,7 @@ Explore the STGraph documentation and tutorials to get started with writing and
    :caption: Tutorials
    :glob:
 
-   tutorials/gnn
+   tutorials/gcn_cora
 
 .. toctree::
    :maxdepth: 1

docs/source/tutorials/gcn_cora.rst

@@ -11,6 +11,8 @@ In this introductory tutorial, you will be able to
 2. Load the Cora dataset provided by STGraph.
 3. Train and evaluate the GCN model for node classification task on the GPU.
 
+You can find the entire source code for this tutorial under the ``tutorials`` directory in our GitHub `repo <https://github.com/bfGraph/STGraph/tree/main/tutorials>`_
+
 The Task At Hand
 ----------------
 

stgraph/nn/pytorch/static/gcn_conv.py

@@ -150,7 +150,7 @@ def forward(input: Tensor, layers: List[GCNConv], graph: StaticGraph):
         """
         if graph.get_ndata("norm") is None:
             raise KeyError("StaticGraph passed to GCNConv forward pass does not contain 'norm' node data")
-        if (len(graph.get_ndata("norm").shape) != SizeConstants.NODE_NORM_SIZE or
+        if (len(graph.get_ndata("norm").shape) != SizeConstants.NODE_NORM_SIZE.value or
                 graph.get_ndata("norm").shape[1] != 1 or
                 graph.get_ndata("norm").shape[0] != graph.get_num_nodes()):
             raise ValueError("Node data 'norm' passed to GCNConv should be of shape (num_nodes, 1)")

tutorials/gcn/cora/README.md

@@ -0,0 +1,3 @@
+# Cora Publication Prediction using Graph Convolutional Networks (GCN)
+
+You can find the detailed tutorial in the documentation page.

tutorials/gcn/cora/model.py

@@ -1,13 +1,38 @@
+"""Graph Convolutional Network Model."""
+
 from __future__ import annotations
 
-import torch.nn as nn
-import torch.nn.functional as F
+from typing import TYPE_CHECKING
+
+from torch import Tensor, nn
+from torch.nn.functional import relu
 
 from stgraph.nn.pytorch.static.gcn_conv import GCNConv
-from stgraph.graph import StaticGraph
+
+if TYPE_CHECKING:
+    from stgraph.graph import StaticGraph
 
 
 class GCN(nn.Module):
+    r"""Graph Convolutional Network Model.
+
+    A multi-layer Graph Convolutional Network Model for node classification task.
+
+    Parameters
+    ----------
+    graph : StaticGraph
+        The input static graph the GCN model operates.
+    in_feats : int
+        Number of input features.
+    n_hidden : int
+        Number of hidden units in a hidden layer.
+    n_classes : int
+        Number of output classes.
+    n_hidden_layers : int
+        Number of hidden layers.
+
+    """
+
     def __init__(
         self: GCN,
         graph: StaticGraph,
@@ -16,22 +41,36 @@ def __init__(
         n_classes: int,
         n_hidden_layers: int,
     ) -> None:
-        super(GCN, self).__init__()
+        r"""Graph Convolutional Network Model."""
+        super().__init__()
 
         self._graph = graph
         self._layers = nn.ModuleList()
 
         # input layer
-        self._layers.append(GCNConv(in_feats, n_hidden, F.relu, bias=True))
+        self._layers.append(GCNConv(in_feats, n_hidden, relu, bias=True))
 
         # hidden layers
-        for i in range(n_hidden_layers):
-            self._layers.append(GCNConv(n_hidden, n_hidden, F.relu, bias=True))
+        for _ in range(n_hidden_layers):
+            self._layers.append(GCNConv(n_hidden, n_hidden, relu, bias=True))
 
         # output layer
         self._layers.append(GCNConv(n_hidden, n_classes, None, bias=True))
 
-    def forward(self: GCN, features):
+    def forward(self: GCN, features: Tensor) -> Tensor:
+        r"""Forward pass of the GCN model.
+
+        Parameters
+        ----------
+        features : Tensor
+            Input features for each node in the graph.
+
+        Returns
+        -------
+        Tensor :
+            The output features after applying all the GCN layers.
+
+        """
         h = features
         for layer in self._layers:
             h = layer.forward(self._graph, h)

tutorials/gcn/cora/train.py

@@ -1,45 +1,64 @@
+r"""Script to train GCN on Cora dataset."""
+
+import sys
 import traceback
 
 import torch
-import torch.nn.functional as F
-
-from stgraph.benchmark_tools.table import BenchmarkTable
-from stgraph.dataset import CoraDataLoader
-from stgraph.graph.static.static_graph import StaticGraph
 from model import GCN
+from torch.nn.functional import cross_entropy
 from utils import (
     accuracy,
     generate_test_mask,
     generate_train_mask,
-    row_normalize_feature,
     get_node_norms,
+    row_normalize_feature,
 )
 
+from stgraph.benchmark_tools.table import BenchmarkTable
+from stgraph.dataset import CoraDataLoader
+from stgraph.graph.static.static_graph import StaticGraph
+
 
 def train(
     lr: float,
     num_epochs: int,
     num_hidden: int,
     num_hidden_layers: int,
-    weight_decay: float
-):
+    weight_decay: float,
+) -> None:
+    r"""Script to train GCN on Cora dataset.
+
+    Parameters
+    ----------
+    lr : float
+        Learning Rate.
+    num_epochs : int
+        Number of Epochs.
+    num_hidden : int
+        Number of hidden units in hidden layer.
+    num_hidden_layers : int
+        Number of hidden layers.
+    weight_decay : float
+        Weight decay value for L2 regularization.
+
+    """
     if not torch.cuda.is_available():
         print("CUDA is not available")
-        exit(1)
+        sys.exit(1)
 
     cora = CoraDataLoader()
 
     node_features = row_normalize_feature(
-        torch.FloatTensor(cora.get_all_features())
+        torch.FloatTensor(cora.get_all_features()),
     )
     node_labels = torch.LongTensor(cora.get_all_targets())
     edge_weights = [1 for _ in range(cora.gdata["num_edges"])]
 
     train_mask = torch.BoolTensor(
-        generate_train_mask(cora.gdata["num_nodes"], 0.7)
+        generate_train_mask(cora.gdata["num_nodes"], 0.7),
     )
     test_mask = torch.BoolTensor(
-        generate_test_mask(cora.gdata["num_nodes"], 0.7)
+        generate_test_mask(cora.gdata["num_nodes"], 0.7),
     )
 
     torch.cuda.set_device(0)
@@ -51,7 +70,7 @@ def train(
     cora_graph = StaticGraph(
         edge_list=cora.get_edges(),
         edge_weights=edge_weights,
-        num_nodes=cora.gdata["num_nodes"]
+        num_nodes=cora.gdata["num_nodes"],
     )
 
     cora_graph.set_ndata("norm", get_node_norms(cora_graph))
@@ -61,16 +80,16 @@ def train(
         in_feats=cora.gdata["num_feats"],
         n_hidden=num_hidden,
         n_classes=cora.gdata["num_classes"],
-        n_hidden_layers=num_hidden_layers
+        n_hidden_layers=num_hidden_layers,
     ).cuda()
 
-    loss_function = F.cross_entropy
+    loss_function = cross_entropy
     optimizer = torch.optim.Adam(
-        model.parameters(), lr=lr, weight_decay=weight_decay
+        model.parameters(), lr=lr, weight_decay=weight_decay,
     )
 
     table = BenchmarkTable(
-        f"STGraph GCN on CORA dataset",
+        "STGraph GCN on CORA dataset",
         ["Epoch", "Train Accuracy %", "Loss"],
     )
 
@@ -91,7 +110,7 @@ def train(
             train_acc = accuracy(logits[train_mask], node_labels[train_mask])
 
             table.add_row(
-                [epoch, float(f"{train_acc * 100:.2f}"), float(f"{loss.item():.5f}")]
+                [epoch, float(f"{train_acc * 100:.2f}"), float(f"{loss.item():.5f}")],
             )
         print("Training Ended")
         table.display()