Different number of aggretation layers for each edge_type in HeteroData()

[Jerome-Cong]

I have a HeteroData() object including two types of nodes named place and room and three types of edges place with place, place sampled_by room and place connected_by room .
The thing I want to do is to first do 3 layers of aggregation for the first two types of edges and then do one layer for the last one. Thus I first try the following structure example:

import torch
from torch_geometric.nn import GCNConv, GATv2Conv
# Initializa the GNN
in_features = 6
graph_conv_features = [64, 64, 64]
local_convs = torch.nn.ModuleList()
global_convs = torch.nn.ModuleList()
local_convs.append(HeteroConv({
            ('places', 'with', 'places'): GCNConv(in_features, graph_conv_features[0]),
            ('places', 'nearby_to', 'room'): GATv2Conv(in_features, graph_conv_features[0], add_self_loops=False),
        }, aggr='sum'))
global_convs.append(HeteroConv({
            ('places', 'connected_by', 'room'): TransformerConv(graph_conv_features[-1], graph_conv_features[-1]),
        }, aggr='sum'))
for i in range(1, len(graph_conv_features)-1):
            local_convs.append(HeteroConv({
                ('places', 'with', 'places'): GCNConv(graph_conv_features[i - 1], graph_conv_features[i]),
                ('places', 'nearby_to', 'room'): GATv2Conv(graph_conv_features[i-1], graph_conv_features[i],
                                                           add_self_loops=False),
            }, aggr='sum'))
# Forward, `x_dict` is the node feature dict of the `HeteroData()` object
for conv in local_convs:
            x_dict = conv(x_dict, edge_index_dict)
            x_dict = {key: h.relu() for key, h in x_dict.items()}
for conv in global_convs:
            x_dict = conv(x_dict, edge_index_dict)
            x_dict = {key: h.relu() for key, h in x_dict.items()}

However I found that after forward the global_convs, the node feature dict only has one key named room. Does this mean that HeteroConv() does not allow only do message passing with one edge type?

[wsad1]

global_convs only has a single GNN with message going from places to room, so it returns a dict with only room. The problem in the above code is that you are overwriting x_dict. Instead try something like

for conv in local_convs:
            x_dict.update(conv(x_dict, edge_index_dict))
            x_dict = {key: h.relu() for key, h in x_dict.items()}
for conv in global_convs:
            x_dict.update(conv(x_dict, edge_index_dict))
            x_dict = {key: h.relu() for key, h in x_dict.items()}