Question about qm9_nn_conv.py in examples

[hieuth133]

Hi,
I have a question about GRU and Set2Set in the qm9_nn_conv.py in examples. What are their purposes here?
Here is the model:

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.lin0 = torch.nn.Linear(dataset.num_features, dim)

        nn = Sequential(Linear(5, 128), ReLU(), Linear(128, dim * dim))
        self.conv = NNConv(dim, dim, nn, aggr='mean')
        self.gru = GRU(dim, dim)

        self.set2set = Set2Set(dim, processing_steps=3)
        self.lin1 = torch.nn.Linear(2 * dim, dim)
        self.lin2 = torch.nn.Linear(dim, 1)

    def forward(self, data):
        out = F.relu(self.lin0(data.x))
        h = out.unsqueeze(0)

        for i in range(3):
            m = F.relu(self.conv(out, data.edge_index, data.edge_attr))
            out, h = self.gru(m.unsqueeze(0), h)
            out = out.squeeze(0)

        out = self.set2set(out, data.batch)
        out = F.relu(self.lin1(out))
        out = self.lin2(out)
        return out.view(-1)

[rusty1s]

This is further specified in the Neural Message Passing for Quantum Chemistry paper (Section 5). In particular, the proposed architecture includes:

1. weight tying at each iteration (using the same NNConv for a fixed number of steps)
2. Using GRU as the update function, which helps in maintaining node representations from earlier layers in the final node-level embeddings.
3. Set2Set is used as the readout function to obtain graph-level representations from node-level ones.

Hope this clarifies your issues :)