Graph level regression task with similar homogenous graphs

[rezabrizi]

Dear PyG community,

I am using GATConv for a regression task on cascade graphs to predict final number of activations in cascade graphs. Cascade graphs are made on the same underlying graph to produce a data point such as Data(x=[4039, 4], edge_index=[2, 176468], y=[1], cascade_name='0'). Node features are [seed_activation_flag, centrality 1, centrality 2, centrality 3]. seed_activation_flag determines if the node is a seed in a given cascade. As you can see centrality remains the same for all cascade graphs since they are static features. I use 8 layers as the underlying graph has a diameter of 8. I am trying to use a deep NN to predict final activation count. Is this a correct way to perform a graph level regression task?

Also, currently my model outputs in the shape torch.Size([1, 1]) but the data.y.shape is torch.Size([1]). Should I worry about the shape?

import torch
from torch_geometric.nn import GATConv, global_add_pool
import torch.nn.functional as F
import torch.nn as nn

class GAT(torch.nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
        super(GAT, self).__init__()
        # mlp for encoding
        self.mlp1 = nn.Linear(input_dim, hidden_dim)
        self.hidden_gat_layers = nn.ModuleList([GATConv(hidden_dim, hidden_dim, heads=1, concat=False, dropout=0.2) for _ in range(num_layers)])
        # mlp for decoding
        self.mlp2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, data):   
        x, edge_index, batch = data.x, data.edge_index, data.batch
        # encode
        x = self.mlp1(x)
        for i, layer in enumerate(self.hidden_gat_layers): 
            x_in = x
            x = F.relu(layer(x, edge_index))
            # Residual Connection
            x = x + x_in 
        x = global_add_pool(x, batch)
        # decode
        x = self.mlp2(x) 
        return x

[rusty1s]

Your model architecture looks fine. If this is a binary classification task, you probably need to reshape the output via .view(-1).