Graph classification with node and edge features benchmark

[mdanb]

Can someone point me to an example of a benchmark graph classification dataset where both the nodes and edges have features? What is the SOTA GNN for it? I browsed TUDatasets, but I could hardly find any that have both edge and node features. For the few that include both, like BZR_MD, I couldn't find the GNN benchmark performance for it. The reason I'm asking is because I constructed the following model (see below) which uses GINE for my dataset which includes both node and edge features. The model always ends up predicting the majority class in my dataset for all test examples (even though I don’t have much class imbalance). I'm unsure if the problem is because of the data, or because there something wrong with my implementation. So I thought I'd test out my model on some benchmark in order to see how it performs on it, which would help me figure out if the problem is mostly from the data or from the implementation of the model. Of course, looking at the code and confirming that my implementation below is not "buggy" would help too. I added in comments in case you can help out with this!

import torch
import torch.nn.functional as F
from torch_geometric.nn import GINEConv
from torch_geometric.nn import global_add_pool, global_mean_pool

class GNN(torch.nn.Module):
    def __init__(self, input_dim, output_dim, config, edge_attr, edge_index, device):  # I pass in edge_attr and edge_index because they are the same for all graphs i.e the samples (graphs) only differ because of the node features
        super(GNN, self).__init__()
        self.pooling_method = config['pooling_method'] # config is just a dictionary. The names within the brackets are all mostly self explanatory. I added an explanation otherwise.
        hidden_dim = config['hidden_layer_size']
        self.num_layers = config['num_layers']  # num convolution layers. Does not include MLP layers which come after message passing. 
        self.dropout = config['dropout_rate']
        self.convs = torch.nn.ModuleList()
        self.conv_type = config['conv_type'] # the convolution layer e.g GINE versus GCN, but right now only supports GINE
        self.consider_graph_level_feature = config['consider_graph_level_feature']
        self.edge_attr = edge_attr
        self.edge_index = edge_index.to(device)
        self.device = device

        self.convs.append(self.build_conv_model(input_dim, hidden_dim, self.conv_type))  # See build_conv_model() definition below
        if config['conv_type'] == 'GINE':
            self.edge_lins = torch.nn.ModuleList()                          # edge_lins holds the linear layers used to process edge attributes, since GINE adds edge embeddings to node embeddings, hence their dimensions must match up
            self.edge_lins.append(torch.nn.Linear(input_dim, hidden_dim))

        self.node_layer_norms = torch.nn.ModuleList()  # Layer normalization for nodes
        self.edge_layer_norms = torch.nn.ModuleList()  # Layer normalization for edges

        for _ in range(self.num_layers - 1):
            self.node_layer_norms.append(torch.nn.LayerNorm(hidden_dim))
            self.edge_layer_norms.append(torch.nn.LayerNorm(hidden_dim))
            self.convs.append(self.build_conv_model(hidden_dim, hidden_dim, self.conv_type))
            if config['conv_type'] == 'GINE':
                self.edge_lins.append(torch.nn.Linear(hidden_dim, hidden_dim))

        # post-message-passing
        if config['consider_graph_level_feature']:
            self.post_mp = torch.nn.Sequential(
                torch.nn.Linear(hidden_dim+1, hidden_dim), torch.nn.ReLU(), torch.nn.Dropout(self.dropout),  # Adds one because graph level feature is 1D (only 1 graph level feature is available for my data)
                torch.nn.Linear(hidden_dim, output_dim), torch.nn.ReLU()) 
        else:
            self.post_mp = torch.nn.Sequential(
                torch.nn.Linear(hidden_dim, hidden_dim), torch.nn.ReLU(), torch.nn.Dropout(self.dropout),
                torch.nn.Linear(hidden_dim, output_dim), torch.nn.ReLU())

    def build_conv_model(self, input_dim, hidden_dim, conv_type):
        if (conv_type == 'GINE'):
            return GINEConv(torch.nn.Sequential(torch.nn.Linear(input_dim, hidden_dim),
                                  torch.nn.ReLU(), torch.nn.Linear(hidden_dim, hidden_dim)))

    def forward(self, data):
        x, batch = data.x, data.batch
        edge_attr = self.edge_attr
        edge_index = self.edge_index

        for i in range(self.num_layers):
            x = self.convs[i](x, edge_index, edge_attr)   # It is not necessary to process edge_attr first before doing this operation because both x and edge_attr start out as having 1 feature only. So edge_attr is basically like edge weights, and x like 1D node features
            x = F.dropout(x, p=self.dropout, training=self.training)
            if self.conv_type == 'GINE':
                edge_attr = self.edge_lins[i](edge_attr)
                edge_attr = F.relu(edge_attr)
                edge_attr = F.dropout(edge_attr, p=self.dropout, training=self.training)

            if not i == self.num_layers - 1:
                x = self.node_layer_norms[i](x)
                if self.conv_type == 'GINE':
                    edge_attr = self.edge_layer_norms[i](edge_attr)

        if (self.pooling_method == "mean"):
            x = global_mean_pool(x, batch)
        elif (self.pooling_method == "add"):
            x = global_add_pool(x, batch)

        if self.consider_graph_level_feature:
            x = self.post_mp(torch.cat((x, torch.tensor(data.u).float().to(self.device)), dim=1))  # data.u is the graph level feature
        else:
            x = self.post_mp(x)

        return F.log_softmax(x, dim=1)

    def loss(self, pred, label):
        return F.nll_loss(pred, label)

[rusty1s]

It's true that the TUDatasets do not really have great edge feature support. I guess one particular use-case where edge features are important is molecular learning, and you can evaluate your model on the ZINC dataset or the datasets provided by OGB.

[mdanb]

OGB is great, thanks for pointing it out!

[HanMihan232]

hi mdanb, If you found a benchmark in ogb that has vertex features and edge features, can you please share?