GINConv for heterogeneous graph

[arzaatri]

I'm attempting to use GINConv for heterogeneous graph prediction. I have set up an initial network like so:

num_type1_feats = 12
num_type2_feats = 3
hidden_dim = 32

ginconv_mlp_1 = nn.Sequential(
    nn.Linear(num_type1_feats, hidden_dim),
    nn.BatchNorm1d(hidden_dim),
    nn.ReLU()
)

ginconv_mlp_2 = nn.Sequential(
    nn.Linear(num_type2_feats, hidden_dim),
    nn.BatchNorm1d(hidden_dim),
    nn.ReLU()
)

class HeteroGNN4(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels, num_layers):
        super().__init__()

        self.convs = torch.nn.ModuleList()
        for _ in range(num_layers):
            conv = HeteroConv({
                ('type1', 'contains', 'type1'): GINConv(ginconv_mlp_1),
                ('type1', 'contains', 'type2'): GINConv(ginconv_mlp_2)
            }, aggr='sum')
            self.convs.append(conv)

        self.readout = SoftmaxAggregation(learn=True)
        self.lin = Linear(hidden_channels, out_channels)

    def forward(self, x_dict, edge_index_dict, batch_dict):

        for conv in self.convs:
            x_dict = conv(x_dict, edge_index_dict)
            x_dict = {key: x.relu() for key, x in x_dict.items()}

        
        x = [self.readout(x_dict[v], batch_dict[v]) for v in x_dict.keys()]
        x = sum(x)
        return relu(self.lin(x))

I load a HeteroData batch and feed it into the network. Here's the error I see:

RuntimeError                              Traceback (most recent call last)
File <command-2780115717529706>, line 1
----> 1 model(batch.x_dict, batch.edge_index_dict, batch.batch_dict)

File /local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File /local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1525         or _global_backward_pre_hooks or _global_backward_hooks
   1526         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527     return forward_call(*args, **kwargs)
   1529 try:
   1530     result = None

File <command-2780115717529695>, line 43, in HeteroGNN4.forward(self, x_dict, edge_index_dict, batch_dict)
     39 def forward(self, x_dict, edge_index_dict, batch_dict):
     40 
     41     # Convolutions will operate jointly on all node types
     42     for conv in self.convs:
---> 43         x_dict = conv(x_dict, edge_index_dict)
     44         x_dict = {key: x.relu() for key, x in x_dict.items()}
     47     x = [self.readout(x_dict[v], batch_dict[v]) for v in x_dict.keys()]

File /local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File /local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1525         or _global_backward_pre_hooks or _global_backward_hooks
   1526         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527     return forward_call(*args, **kwargs)
   1529 try:
   1530     result = None

File /local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/torch_geometric/nn/conv/hetero_conv.py:159, in HeteroConv.forward(self, *args_dict, **kwargs_dict)
    156     if not has_edge_level_arg:
    157         continue
--> 159     out = self.convs[edge_type](*args, **kwargs)
    160     out_dict[dst].append(out)
    162 for key, value in out_dict.items():

File /local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File /local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1525         or _global_backward_pre_hooks or _global_backward_hooks
   1526         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527     return forward_call(*args, **kwargs)
   1529 try:
   1530     result = None

File /local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/torch_geometric/nn/conv/gin_conv.py:84, in GINConv.forward(self, x, edge_index, size)
     82 x_r = x[1]
     83 if x_r is not None:
---> 84     out = out + (1 + self.eps) * x_r
     86 return self.nn(out)

RuntimeError: The size of tensor a (12) must match the size of tensor b (3) at non-singleton dimension 1

Notably, 12 and 3 are the number of node features for node type 1 and node type 2, respectively.

The error does not occur if I change the GNN operator for ('type1', 'contains', 'type2'): to GATConv((-1, -1), hidden_channels, add_self_loops=False)

Could anyone give guidance as to why this error happens and how to fix it?

EDIT: This also happens when using an out-of-the-box GIN model, like so:

model = GIN(-1, 32, 4, 1)
node_types = ['type1', 'type2']
edge_types = [('type1', 'contains', 'type1'), ('type1', 'contains', 'type2')]
metadata = (node_types, edge_types)

model = to_hetero(model, metadata)

[rusty1s]

GINConv cannot handle different feature dimensionalities across source and destination nodes. As such, you might want to consider to first project node features of all node types into the same space.

[Shi-Jialiang]

So can we have a version of GINConv for unidirectional bipartite graph like GAT🤔？

[rusty1s]

You can by first projecting all node types to the same feature dimensionality.

[Shi-Jialiang]

emmm...sorry these days I am using DGL, and they have an api called HeteroGraphConv to manage the Conv by this format for heterogeneous graph:

import dgl.nn.pytorch as dglnn

conv = dglnn.HeteroGraphConv({
  'follows' : dglnn.GraphConv(...),
  'plays' : dglnn.GraphConv(...),
  'sells' : dglnn.SAGEConv(...)},
  aggregate='sum')

But they don't offer the GINConv for the single unidirectional bipartite graph (Here, like the relation 'follows' in the example) like GraphConv or SAGEConv. So, do you think it's possible to use the GINConv to all the specific relation graph and then aggregate the result? At least I think in the dgl's situation, I should correct the source code. Thank you so much for noticing my question!

[arzaatri]

What I did was this:

• I had two node types A and B
• A has 10 features, B has 5. So 15 total between the two
• I transformed the 10-element feature tensor for A to a 15-element vector, whose elements 0:10 were the features for A. The rest were 0
• Likewise for nodes of type B, but zeroing elements 0:10 and putting their features in elements 10:15

Someone can correct me if this approach has issues

[rusty1s]

GINConv is well supported by HeteroConv - it's only requirement is that feature dimensionality is shared. You can easily do

x_dict = {key: self.lin_dict[key](x) for key, x in x_dict.items()}
x_dict = conv(x_dict, edge_index_dict)

where conv is defined as

conv = HeteroConv({
    edge_type1: GINConv(...),
    edge_type2: GINConv(...),
}