Support for heterogeneous graphs

❓ Questions & Help

I new to gnn and was wondering if PyG has any support for modelling heterogeneous graphs. I’ve seen a HinSAGE and metapath2vec in the Stellar library but I am a PyTorch diva! Any plans to integrate these? Or, maybe you could describe how I can maybe use existing blocks to create these?

Also, is there support for the dataloading process? DGL has heterogeneous graphs written on top of networkx. Is there any easy way to port these data structures over into PyG.

Any feedback helps!

[fgerzer]

To the best of my knowledge, pytorch-geometric does not have any specific prebuilt ways of dealing with heterogeneous graphs.

That being said, it’s still fairly easy to implement, and I’ve done so previously on a power grid dataset. I’ll use a simplified form of this as an example.

From memory, what you have to do is:

1. Create a corresponding Data object

   • This simply means not just using Data’s x, edge_index, and edge_attr attributes, but adding whatever attributes you need.
   • These will automatically be batched (the corresponding code if you define it correspondingly in your Data object. For example,

   class PowerData(Data):
       def __inc__(self, key, value):
           increasing_funcs = ["load_to_bus", "branch_index", "gen_to_bus"]
           if key in increasing_funcs:
               return len(self["bus"])
           else:
               return 0
   

   ensures that the keys in increasing_funcs will be treated just like edge_index, i.e. on batching the index will increase. For example, two branch_index (my equivalent to edge_index) of [0, 1] , [1, 0] and [0, 1], [1, 0] will then be automatically batched to [0, 1, 2, 3], [1, 0, 3, 2]. (For me, this was fairly simple because load_to_bus and gen_to_bus were one-dimensional. For you, you might have to add different values to different dimensions; maybe something like return np.array([len(self["gen"]), len(self["bus"]])

   • My Data objects had bus, generator and load features, and bus_to_bus, gen_to_bus, and load_to_bus edges (with the same format as edge_index). Additionally, there were branch features (related to the bus_to_bus edges).

2. Create a model to deal with heterogeneous graphs

   • This might be as simple as having one encoding model for each node type, then concatenating node features and edge indices, or might be as complicated as having different GNNs for each of the node interactions. For example, one layer that I evaluated looked as follows:

    def forward(self, data):
        bus = data.bus
        gen = self.subnets["gen"](torch.cat([data.gen, bus[data.gen_to_bus]], dim=-1))
        load = self.subnets["load"](torch.cat([data.load, bus[data.load_to_bus]], dim=-1))
        bus = self.subnets["bus"](torch.cat([
            bus,
            self._scatter_items(gen, data.gen_to_bus, bus.shape[0]),
            self._scatter_items(load, data.load_to_bus, bus.shape[0]),
        ], dim=-1))
        src, dest = data.branch_index
        branch = self.subnets["branch"](torch.cat([bus[src], data.branch_attr, bus[dest]], dim=-1))
        bus_neighbours = self.subnets["bus_and_branch"](torch.cat([bus[dest], data.branch_attr], dim=-1))
        bus_neighbours = scatter_add(bus_neighbours, src, dim=0, dim_size=bus.shape[0])
        bus = self.subnets["bus_and_neighbours"](torch.cat([bus_neighbours, bus], dim=-1))
        data.bus = bus
        data.gen = gen
        data.load = load
        data.branch_attr = branch
        return data

[JiaxuanYou]

Heterogeneous graph learning is supported in PyG 2.0: https://pytorch-geometric.readthedocs.io/en/latest/notes/heterogeneous.html