Zinc Dataset Example using PNAConv

[arake001]

I was going through the examples of PyTorch geometric for the ZINC dataset and found a code using PNAConv, which uses edge_embedding and node_embedding. The choice of embedding sizes, specifically node_embedding (21, 75) and edge_embedding (4, 50), appears intriguing. I'm curious about the rationale behind these selections.

from torch_geometric.datasets import ZINC
dataset = ZINC(path)
print(dataset.data)

Data(x=[5094322, 1], edge_index=[2, 10959434], edge_attr=[10959434], y=[220011])

Example code:

class Net(torch.nn.Module):
def init(self):
super().init()

    self.node_emb = Embedding(21, 75)
    self.edge_emb = Embedding(4, 50)

    aggregators = ['mean', 'min', 'max', 'std']
    scalers = ['identity', 'amplification', 'attenuation']

    self.convs = ModuleList()
    self.batch_norms = ModuleList()
    for _ in range(4):
        conv = PNAConv(in_channels=75, out_channels=75,
                       aggregators=aggregators, scalers=scalers, deg=deg,
                       edge_dim=50, towers=5, pre_layers=1, post_layers=1,
                       divide_input=False)
        self.convs.append(conv)
        self.batch_norms.append(BatchNorm(75))

    self.mlp = Sequential(Linear(75, 50), ReLU(), Linear(50, 25), ReLU(),
                          Linear(25, 1))

def forward(self, x, edge_index, edge_attr, batch):
    x = self.node_emb(x.squeeze())
    edge_attr = self.edge_emb(edge_attr)

    for conv, batch_norm in zip(self.convs, self.batch_norms):
        x = F.relu(batch_norm(conv(x, edge_index, edge_attr)))

    x = global_add_pool(x, batch)
    return self.mlp(x)

[rusty1s]

These are just the maximum indices present in data.x and data.edge_attr.

[arake001]

Okay. Understood. Thank you for your help! Also, is there any example code for the dataset AQSOL?