Handle shape of node representations

[afonso-sousa]

Hi. I am using GNNs to get node representations to feed to downstream models. I have the following model:

class GCNEncoder(nn.Module):
    def __init__(self, num_features, hidden_size, dropout=0):
        super().__init__()
        self.conv1 = GCNConv(num_features, hidden_size)
        self.conv2 = GCNConv(hidden_size, hidden_size)
        self.relu = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=dropout)
        self.lin = nn.Linear(hidden_size, 50265)

    def forward(self, data, emb_weights):
        x, edge_index = data.x, data.edge_index
        x = (x - x.mean(dim=0, keepdim=True)) / x.std(dim=0, keepdim=True)

        x = self.conv1(x, edge_index)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.conv2(x, edge_index)
        x = self.relu(x)

        # x = self.lin(x) # [722, 50265]
        q = F.linear(x, emb_weights) # 862 x 50265

        x = global_add_pool(
            x, data.batch, size=data.num_graphs
        )  # batch_size x hidden_size => [32, 256]

        return x

I want my network to output a representation of each node, assuming there can be 50265 of them (num_nodes). Specifically, I want a way to produce a representation of shape batch_size x num_nodes x hidden_size and hidden_size x num_nodes. Respectively, I can get batch-wise representations of shape batch_size x hidden_size using global_add_pool but do not know how to add the middle dimension of size num_nodes; and can get a two-dimensional representation with a linear, but the first dimension is of size 862 instead of hidden (512), and do not know if this q will learn what I want since I am using weights from an Embedding layer (emb_weights) hence why I need the node representations shaped as hidden_size x num_nodes.
Thanks in advance for any help you can provide.

[rusty1s]

You want to drop the global_add_pool here. In addition, your final Linear module should output hidden_size features instead of 50265. Then, you can do:

x = self.lin(x)
x, mask = to_dense_batch(x, batch)
return x  # [batch_size, num_nodes, num_features]