How to construct several neural networks for different node pairs in a model?

[thu-wangz17]

❓ Questions & Help

Hi, I have a question analogous to heterogeneous GNN. For example, I have two kinds of nodes, which are denoted by A and B respectively. I want to construct three same neural network for A-A pairs, A-B pairs and B-B pairs but with different weights, or construct three different neural networks for these pairs. Since the model will be trained with a batch which will consist of several graphs, I don't know how could I implement?

[rusty1s]

Great question! This is currently not that well-implemented, and I am working on a cleaner interface. PyG provides all tools to implement A-A, A-B and B-B message passing, but batch construction is currently quite cumbersome: I suggest to inherit from the Data class and re-implement the __inc__ method, which describes how respective edge indices should be added together:

class MyHetoData(Data):
    def __inc__(self, key, value):
        if key == 'edge_index_A_A':
            return num_nodes_A
        elif key == 'edge_index_B_B':
            return num_nodes_B
        elif key == 'edge_index_A_B':
            return torch.tensor([num_nodesA, num_nodes_B])
        else:
            return 0

[thu-wangz17]

Thanks for replying.