Heterogeneous graph and pooling

[minsikseo-cdl]

Hi, I'm working on a heterogeneous graph like below:

Unlike, the example on PyG document, this graph has the same node attributes, but different edge attributes. As in the figure, most of the nodes in A are connected to each other with edges with 4-d edge attributes (A, a-inner, A). Some of nodes in A are connected by edges having no edge attribute (A, a-outer, A). In B, most of nodes are connected by 3-d attributes edges namely (B, b-inner, B). The nodes in C are not connected to each other, but to the nodes in A with no-attrs-edge (A, a-c, C). Finally, A and B are connected with some nodes in them by no-attrs-edge (A, a-b, B).

After creating a graph with HetaroData class, I made this graph undirected by a transform ToUndirected, so that the graph has reverse of edges above (e.g., (A, rec_a-c, C)).

In this situation, the goal is to embed the entire graph into a vector. I've planed to design a GNN having some convolution type of layers and pooling between them. I am also confused about building a heterogeneous GNN model, but it will be figured out soon. But I cannot determine if it is possible to pool this type of heterogeneous graph or not. So, here is my question.

1. Is it possible to pool a heterogeneous graph like above?
2. If it is, how can I implement this pooling layer? Any hint please?

Thanks.

[rusty1s]

Pooling can be achieved by pooling each node type and combining their graph-level outputs afterwards, e.g.:

loader = DataLoader(dataset, ...)

for data in loader:
    x_dict = {key: global_mean_pool(data[key].x, data[key].batch) for key in data.node_types}
    out = torch.stack(x_dict.values(), dim=0).sum(dim=0)

[minsikseo-cdl]

Oh, thank you so much.
I had to mention that I'd like to pool "locally," not globally (by aggregation), though.
What I want to do is hierarchical pooling between conv layers.
I think, there is no way to directly do this on a heterogeneous graph.
But, I have ideas:

1. Making a graph be homogeneous in the first conv layer by updating edge attributes to have the same dimension. But, how can I update edge attributes in PyG?
2. From the question of the first idea: in my case, since the graph has heterogeneous edge and homogeneous nodes, I can transform the graph into a line graph using pyg.transforms.LineGraph. Then I have a heterogeneous graph in terms of node attributes, and by passing one heterogeneous layer, I can get a homogeneous graph. But, LineGraph seems not working in my case. I've tried:
   I have

print(type(hetero), hetero.num_nodes, hetero.num_edges)
>>> <class 'torch_geometric.data.hetero_data.HeteroData'> 177652 1096526

homo = hetero.to_homogeneous()
print(homo)
>>> Data(x=[177652, 3], edge_index=[2, 1096526], edge_type=[1096526])

and then

line = LineGraph()(homo)
print(line)
>>> Data(x=[177652, 3], edge_index=[2, 67674764], num_nodes=3031928)

I don't get why the number of nodes is not equivalent to the number of edges of the original graph. So I cannot assign the attributes for both nodes and edges.

Could you let me know how can I update the edge attributes? Or, any idea? 😭

[rusty1s]

Oh, that's interesting. I don't think there exists any pooling approach for heterogeneous graphs yet (both in research and in PyG). Here are some thoughts.

1. The question is whether you want to maintain node type and edge type information after pooling. Do you only want to allow pooling of nodes of same node type? If you want to loose the heterogeneous graph information, you might indeed be able to leverage to_homogeneous, and use regular PyG pooling approaches. Edge features need to be updated when creating coarsened graphs, which should be already handled by methods such as torch_geometric.pool.max_pool.
2. Regarding your LineGraph issue, this may be explained by the fact that the graph is directed, not undirected.

[minsikseo-cdl]

I decided to define a custom layer and make a heterogeneous graph homogeneous (same dimensionalities for all node and edge attributes). And then, I'll use diffpool for pooling a homogenized graph. Designing a proper layer setup has to be figured out, but I'm going to try..
Thanks! 😊