From Node Embeddings to Graph Embedding to Linear Regression

[GeorgiosIoannouCoder]

First, I would like to say a huge thank you to the Pytorch Geometric Team and its open Community for all the valuable efforts, time, and useful resources.

Task:
Let us say that we have 3,000 directed torch_geometric.data.Data objects where each node has 10 features and each of the 3,000 Graphs has a different number of nodes (however still each node has 10 features) and each of the 3,000 Graphs has a float Graph-level ground-truth label (i.e. A total of 3,000 float Graph-level ground-truth labels):

1. We want to use GraphSage to produce node embeddings for each of the 3,000 Graphs and increase the features for each node in the 3,000 Graphs from 10 to 128. This will result in increasing the number of features of each node from 10 to 128 and produce the node embeddings.
2. Then, we want to use global mean pooling to generate a graph embedding for each of the 3,000 Graphs. This will result in 3,000 graph embeddings.
3. Finally, we want to use these 3,000 graph embeddings (X) with the 3,000 float Graph-level ground-truth labels (y) in a Linear Regression model.

Questions:

1. Is the following model implementation a good start for the Task?

class GraphSAGE(torch.nn.Module):
    def __init__(
        self,
        in_channels,
        hidden_channels,
        num_layers,
        out_channels,
        dropout,
        act,
        aggr,
        normalize,
        act_first=False,
        act_kwargs=None,
        norm=None,
        norm_kwargs=None,
        jk=None,
        root_weight=True,
        project=False,
        bias=True,
    ):
        super(GraphSAGE, self).__init__()

        torch.manual_seed(12345)

        self.num_layers = num_layers
        self.dropout = dropout
        self.act = act
        self.convs = nn.ModuleList()
        self.act_first = act_first
        self.act_kwargs = act_kwargs
        self.norm = norm
        self.norm_kwargs = norm_kwargs
        self.jk = jk
        self.root_weight = root_weight
        self.project = project
        self.bias = bias

        self.conv1 = SAGEConv(
            in_channels=in_channels,
            out_channels=hidden_channels,
            aggr=aggr,
            normalize=normalize,
        )
        self.conv2 = SAGEConv(
            in_channels=hidden_channels,
            out_channels=hidden_channels,
            aggr=aggr,
            normalize=normalize,
        )
        self.conv3 = SAGEConv(
            in_channels=hidden_channels,
            out_channels=out_channels,
            aggr=aggr,
            normalize=normalize,
        )

        self.lin1 = nn.Linear(hidden_channels, out_channels)

    def forward(
        self,
        x,
        edge_index,
        nodes,
        batch=None,
    ):
        # 1. Obtain node embeddings.
        # Embed each node by performing multiple rounds of message passing.

        h = self.conv1(x, edge_index)
        h = h.relu()
        h = self.conv2(h, edge_index)
        h = h.relu()
        h = self.conv3(h, edge_index)  # Final GNN embedding space.

        # 2. Readout layer.

        batch = torch.zeros(nodes, dtype=int) if batch is None else batch
        h = global_mean_pool(x=h, batch=batch)

        # 3. Apply a final classifier.

        h = F.dropout(input=h, p=self.dropout, training=self.training)
        h = self.lin1(h)

        return h
    

model = GraphSAGE(
    in_channels=10
    hidden_channels=128
    num_layers=3,  # Number of message passing layers.
    out_channels=128,
    dropout=0.2,
    act="relu",
    aggr="mean",
    normalize=True,
)

2. Do we need to train the above model for a certain number of epochs as done in the Graph Classification with Graph Neural Networks Pytorch Geometric Example? We only need to use GraphSage as an intermediate step in our process for the node embeddings. Therefore, do we still need to train for a certain number of epochs or we just call it for one epoch (one time) without the need to update any parameters.
3. We can do our Task in Batches right?
4. Can NeighborLoader be used in this Task?

Thank you very much in advance. I appreciate a lot.

[rusty1s]

1. Yes, this looks good, but there is probably no reason to out-source classification to a separate model. I suggest to train this end-to-end as shown here: https://github.com/pyg-team/pytorch_geometric/blob/master/examples/mutag_gin.py
2. Yes, you will need to train your model. You can do it in a supervised fashion end-to-end (recommended), or train the GNN in an unsupervised fashion to use the embeddings downstream (see examples/graph_sage_unsup.py)
3. Yes, this will be achieved through the DataLoader
4. No, the NeighborLoader is only useful for operating on single large graphs. Use DataLoader instead :)