Notes on Heterogenous Graphs - PyTorch Geometric

In this guide I will share my insights that I gained while working with heterogenous graphs in PyTorch Geometric

Introduction

PyTorch Geometric offers a lot of useful tools to work with graphs. However, the documentation is not always very clear and it can be hard to find the right information. In this post I want to share my insights that I gained while working with heterogenous graphs in PyTorch Geometric. I will try to keep it as short as possible and only focus on the most important aspects. If you want to learn more about the topic I recommend to read the official documentation.

Heterogenous Graphs

A heterogenous graph is a graph that contains different types of nodes and edges. For example, a citation network contains two types of nodes: papers and authors. The edges in this graph can be of different types as well. For example, a paper can cite another paper or an author can write a paper. In this case we have two different types of edges: paper-paper and author-paper.

In my particular usecase I was faced with the challenge of having a graph which contains a robot and objects located in the scene. This implies different features per node. Where the first kind of nodes are robot links and the second kind of nodes are objects. The edges in this graph are of different types as well. For example, a robot link can be connected to another robot link or an object can be connected to a robot link. In this case we have two different types of edges: robot-robot and robot-object.

Data

In PyTorch Geometric a regular graph is represented by a torch_geometric.data.Data object. This data object can contain only one type of nodes and edges. In order to represent a heterogenous graph we need to use a torch_geometric.data.HeteroData object. This object can contain multiple types of nodes and edges. The following code snippet shows how to create a heterogenous graph:

from torch_geometric.data import HeteroData

data = HeteroData()

# Robot link features (node features)
data['robot'].x = torch.randn(5, 16)

# Object features (node features)
data['object'].x = torch.randn(10, 32)

# Robot-robot edges
data['robot', 'connects', 'robot'].edge_index = torch.tensor([
    [0, 1, 1, 2, 3, 4],
    [1, 0, 2, 1, 4, 3],
])

# Robot-object edges undirected
data['robot', 'relates', 'object'].edge_index = torch.tensor([
    [0, 0, 1, 1, 2, 2, 3, 3, 4, 4],
    [0, 1, 0, 1, 2, 3, 2, 3, 4, 3],
])

In the example we can see a heterogenous graph with two types of nodes: robot and object. There are 5 robot links and 10 objects. The robot links have 16 features and the objects have 32 features. Edges betweeen the nodes are defined by the triplets (source, relation, target). In this case we have two types of edges: robot-robot and robot-object. The robot-robot edges are defined by the triplet (robot, connects, robot). The robot-object edges are defined by the triplets (robot, relates, object) and (object, relates, robot). It is important to note, that the first column of the edge index always corresponds to the source nodes and the second column to the target nodes. Additionally it’s required that the node indices of each node type start with 0. This means there is no global node index.

Message Passing and Models

For information on message passing and modules please refer to the official documentation.

Lazy Initialization of input nodes

In the documentation a technique called lazy initilization is introduced. This is neccessary since the input nodes are not known in advanace. In the following example we can see how to use this technique:

class GAT(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = pyg_nn.GATConv((-1, -1), hidden_channels, add_self_loops=False)
        self.lin1 = pyg_nn.Linear(-1, hidden_channels)
        self.conv2 = pyg_nn.GATConv((-1, -1), out_channels, add_self_loops=False)
        self.lin2 = pyg_nn.Linear(-1, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index) + self.lin1(x)
        x = x.relu()
        x = self.conv2(x, edge_index) + self.lin2(x)
        return x

The input features are represented by the tuple (-1, -1). This represents a two dimensional input tensor. For one dimensional input feature we would use just a single -1 in the corresponding function call.

To initialize the input features we can use the following code snippet:

with torch.no_grad():  # Initialize lazy modules.
        init_data = train_loader.dataset[0].to(device)
        out = model(init_data.x_dict, init_data.edge_index_dict)

This will extract the input features from the first data object in the dataset and initialize the lazy modules. The output of the model is not important and can be ignored. But in this call we can also see how we call the model. Instead of just passing a data object to the model we pass a dictionary containing our different node features and edge indices.

Conclusion

In this post I shared my insights that I gained while working with heterogenous graphs in PyTorch Geometric. I hope this post was helpful and you learned something new. If you have any questions or feedback feel free to contact me.

Error messages that I encountered and their solutions