Models from torch_geomtric.nn vs manually implemented Models in PyTorch

1

I am trying to use the built-in models from torch_geometric.nn module. I have also seen many tutorial that most of the models are designed manually by extending the torch.nn.Module class.

What is the difference between the following:

from torch_geometric.nn import GCN
dataset = TUDataset(root=data_root, name=data_name)
train_dataset = dataset[:800]
test_dataset = dataset[800 :]
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)
model = GCN(dataset.num_node_features, 64, 3, dataset.num_classes)

and

from torch_geometric.nn import GCNConv, global_mean_pool, Linear
class MyGCN(torch.nn.Module):
    """GCN"""
    def __init__(self, dim_h):
        super(MyGCN, self).__init__()
        self.conv1 = GCNConv(dataset.num_node_features, dim_h)
        self.conv2 = GCNConv(dim_h, dim_h)
        self.conv3 = GCNConv(dim_h, dim_h)
        self.lin = Linear(dim_h, dataset.num_classes)

    def forward(self, x, edge_index, batch):
        # Node embeddings 
        h = self.conv1(x, edge_index)
        h = h.relu()
        h = self.conv2(h, edge_index)
        h = h.relu()
        h = self.conv3(h, edge_index)

        # Graph-level readout
        hG = global_mean_pool(h, batch)

        # Classifier
        h = F.dropout(hG, p=0.5, training=self.training)
        h = self.lin(h)
        
        return hG, F.log_softmax(h, dim=1)

Also, what is the relation between hidden_channels and batch_size? I cannot find something specific that describes the difference between the two.

2

The PyG implementation of GCN seems to define a few additional methods (such as inference) and seems to abstract the actual GCN class by deriving from BasicGNN as seen here.
hidden_channels seems to be used to initialize the self.convs while the batch_size is used during the inference call while iterating the loader to slice the outputs.
@rusty1s can correct me if I’m missing something as he is the core developer.