How to implement graph neural networks in Mathematica?

Question

Graph neural networks are useful for prediction tasks like predicting walks. In python I use DeepGraph typically, but I'm wondering what can be done in the new Version 12.

Here's a nice introduction blog post to the topic: https://towardsdatascience.com/a-gentle-introduction-to-graph-neural-network-basics-deepwalk-and-graphsage-db5d540d50b3.

Answer 1

Background

Graphs can be thought as generalized images where each node can have arbitrary neighbors instead of strict grid neighbors. The simplest graph network is the graph convolutional network (GCN).

In the convolutional network (CNN), each pixel's value is computed from the aggregation of the neighbor pixels in terms of convolution operations. Similarly, each graph node's value is computed from the aggregation of the values of its neighbor nodes.

A convolution operation can be written as $$ H' = \sigma(cov(H,W)) $$

where $H$ is the value of the pixel, $W$ is the shared learnable convolution kernel, $H'$ is the value of pixel value after this convolution operation and $\sigma$ is the non-linear operation.

Similarly, a simple graph convolution can be written as: $$ H' = \sigma(A H W) $$

where $W$ is the shared learnable weight, and $A$ is the adjacency matrix. Conceptually, a node's value comes from the aggregation of all its neighbors.

Construction

A single graph convolutional layer can be defined as

gnn[adj_, transform_, activation_ : True] := Module[{},
  NetGraph[{"fts_transf" -> NetMapOperator[transform], 
    "dot" -> DotLayer[], 
    "adj" -> 
     NetArrayLayer["Array" -> adj, LearningRateMultipliers -> None], 
    If[activation, "activation" -> Ramp, Nothing]}, {"adj" -> "dot", 
    "fts_transf" -> "dot", 
    If[activation, "dot" -> "activation", "dot" -> NetPort["Output"]]}]
  ]

where we used the NetMapOperator to apply the transform to each node, and multiply the results with the adjacency matrix. In the case where transform is just a linear layer, the graph convolutional layer has a structure

gnn[adj, LinearLayer[16], True]

One simple GCN network is then constructed from multiple graph convolutional layers

net[adj_] := 
 NetGraph[{"gnn1" -> gnn[adj, LinearLayer[32], True], 
   "gnn2" -> gnn[adj, LinearLayer[numClasses], False], 
   "softmax" -> SoftmaxLayer[]}, {"gnn1" -> "gnn2" -> "softmax"}, 
  "Input" -> {numNodes, numFeatures}]
net[adj]

Example

We use the Cora dataset as an example to demonstrate the GCN. Cora dataset is a citation network among papers and each paper is represented as a node. The task is to classify each paper into one of the seven classes. Here is a visualization of the graph with its 2708 nodes, colored by the classes.

We can use the above network to perform this classification task. But before training the network, we need to define a cross entropy loss that can take a node mask as an argument. This node mask defines the separation of the training and test nodes.

loss[mask_] := NetGraph[{
   "mask" -> 
    NetArrayLayer["Array" -> Position[mask, 1], 
     LearningRateMultipliers -> None],
   "masking1" -> ExtractLayer[], "masking2" -> ExtractLayer[],
   "ce" -> 
    CrossEntropyLossLayer["Index", 
     "Input" -> {Length[Position[mask, 1]], numClasses}, 
     "Target" -> 140]
   },
  {"mask" -> NetPort["masking1", "Position"], 
   "mask" -> NetPort["masking2", "Position"],
   NetPort["Input"] -> "masking1", NetPort["Target"] -> "masking2",
   {"masking1", "masking2"} -> "ce"
   }, "Input" -> {numNodes, numClasses}, "Target" -> numNodes]

The network can then be trained as

data = {features -> labels};
trained = 
 NetTrain[net, data, LossFunction -> loss[trainMask], 
  , MaxTrainingRounds -> 100]

And the accuracies can be evaluated as

evalAcc[mask_] := 
 Module[{idx = Flatten@Position[mask, 1], 
   res = Flatten[Ordering[#, -1] & /@ trained[features]]}, 
  Equal @@@ Transpose[{res[[idx]], labels[[idx]]}] // Counts]
testAcc = evalAcc[testMask]
testAcc[True]/(testAcc[True] + testAcc[False]) // N

(*<|False -> 202, True -> 798|>*)
(*0.798*)

We see that the test accuracy is about 80%.

In order to verify that the graph structure indeed helps with the classification of the nodes, we can clear out the graph structure by setting the adjacency matrix to the identity matrix.

trained = 
 NetTrain[net[IdentityMatrix[numNodes]], data, 
  LossFunction -> loss[trainMask], BatchSize -> 1, 
  MaxTrainingRounds -> 100, TargetDevice -> "GPU"]

We can see now the test accuracy decreased to about 50%.

testAcc = evalAcc[testMask]
testAcc[True]/(testAcc[True] + testAcc[False]) // N
(*<|False -> 490, True -> 510|>*)
(*0.51*)

Appendix

Data used in the above example can be obtained in python using the following code

import numpy as np
import torch
import torch_geometric

cora_dataset = torch_geometric.datasets.Planetoid(root="/tmp/data",
name="Cora")[0]

num_nodes = cora_dataset.num_nodes
num_features = cora_dataset.num_features
num_classes = 7

node_features = cora_dataset.x

train_mask = cora_dataset.train_mask
test_mask = cora_dataset.test_mask
val_mask = cora_dataset.val_mask

labels = cora_dataset.y
logits = torch.zeros([num_nodes,num_classes])
for i,pos in enumerate(labels):
  logits[i][pos] = 1

adj_matrix = torch.zeros([num_nodes, num_nodes])
for i,j in cora_dataset.edge_index.T.numpy():
  adj_matrix[i,j] = 1
  adj_matrix[j,i] = 1
for i in range(num_nodes):
  adj_matrix[i,i] = 1

[adj_matrix.numpy(), node_features.numpy(), labels.numpy(), 
  train_mask.numpy().astype(np.int), 
  test_mask.numpy().astype(np.int)]

And

{adj, features, labels, trainMask, testMask} = Normal@%;
{numNodes, numFeatures} = Dimensions[features];
numClasses = Max[labels] + 1;
(*normalize adjacency matrix*)
adj = #/Total[#] & /@ adj;
(*change to one-based index*)
labels = labels + 1;

Reference

1. UvA Deep Learning Tutorials
2. Graph Convolutional Networks
3. Intro to graph neural networks (ML Tech Talks)