Building graph neural networks
Building GNN is as simple as building neural network in Flux. The syntax here is the same as Flux. Chain
is used to stack layers into a GNN. A simple example is shown here:
model = Chain(GCNConv(adj_mat, feat=>h1),
GCNConv(adj_mat, h1=>h2, relu))
In the example above, The first argument adj_mat
is the representation of a graph in form of adjacency matrix. The feature dimension in first layer is mapped from feat
to h1
. In second layer, h1
is then mapped to h2
. Default activation function is given as identity if it is not specified by users.
The initialization function GCNConv(...)
constructs a GCNConv
layer. For most of the layer types in GeometricFlux, a layer can be initialized in at least two ways:
- Initializing with a predefined adjacency matrix or
FeaturedGraph
, followed by the other parameters. For most of the layer types, this is for datasets where each input has the same graph structure. - Initializing without an initial graph argument, only supplying the relevant parameters. This allows the layer to accept different graph structures.
Applying layers
When using GNN layers, the general guidelines are:
- If you pass in a $n \times d$ matrix of node features, and the layer maps node features $\mathbb{R}^d \rightarrow \mathbb{R}^k$ then the output will be in matrix with dimensions $n \times k$. The same ostensibly goes for edge features but as of now no layer type supports outputting new edge features.
- If you pass in a
FeaturedGraph
, the output will be also be aFeaturedGraph
with modified node (and/or edge) features. Addnode_feature
as the following entry in the Flux chain (or simply callnode_feature()
on the output) if you wish to subsequently convert them to matrix form.
Create custom layers
Customizing your own GNN layers are the same as customizing layers in Flux. You may want to reference Flux documentation.