Chapter 3: Building GNN Modules


DGL NN module consists of building blocks for GNN models. An NN module inherits from Pytorch’s NN Module, MXNet Gluon’s NN Block and TensorFlow’s Keras Layer, depending on the DNN framework backend in use. In a DGL NN module, the parameter registration in construction function and tensor operation in forward function are the same with the backend framework. In this way, DGL code can be seamlessly integrated into the backend framework code. The major difference lies in the message passing operations that are unique in DGL.

DGL has integrated many commonly used apinn-pytorch-conv, apinn-pytorch-dense-conv, apinn-pytorch-pooling, and apinn-pytorch-util. We welcome your contribution!

This chapter takes SAGEConv with Pytorch backend as an example to introduce how to build a custom DGL NN Module.