GraphNets.jl

Simple, blazing fast, message-passing graph neural network.

GraphNets.jl is a Julia implementation of DeepMind's Graph Nets and an implementation of the following paper: Relational inductive biases, deep learning, and graph networks

Petar Veličković gives a fantastic explanation of the Graph Nets architecture and graph neural networks in general in the lecture below.

Setup

using GraphNets

X_DE = 10 # Input feature dimension of edges
X_DN = 5 # Input feature dimension of nodes
X_DG = 0 # Input feature dimension of graphs (no graph level input data)
Y_DE = 3 # Output feature dimension of edges
Y_DN = 4 # Output feature dimension of nodes
Y_DG = 5 # Output feature dimension of graphs

block = GNBlock(
    (X_DE,X_DN,X_DG) => (Y_DE,Y_DN,Y_DG)
)

Example 1: Batch of graphs with same structure (same adjacency matrix)

adj_mat = [
    1 0 1;
    1 1 0;
    0 0 1;
] # Adjacency matrix

num_nodes = size(adj_mat, 1)
num_edges = length(filter(isone, adj_mat))

batch_size = 2
edge_features = rand(Float32, X_DE, num_edges, batch_size)
node_features = rand(Float32, X_DN, num_nodes, batch_size)
graph_features = nothing # no graph level input features

x = (
    graphs=adj_mat, # All graphs in this batch have same structure
    ef=edge_features, # (X_DE, num_edges, batch_size)
    nf=node_features, # (X_DN, num_nodes, batch_size)
    gf=graph_features # (X_DG, batch_size)
) |> batch

y = block(x) |> unbatch

@assert size(y.ef) == (Y_DE, num_edges, batch_size)
@assert size(y.nf) == (Y_DN, num_nodes, batch_size)
@assert size(y.gf) == (Y_DG, batch_size)

# Get the output graph edges of the 1st graph
@assert size(y.ef[:,:,1]) == (Y_DE, num_edges)

# Get the output node edges of the 1st graph
@assert size(y.nf[:,:,1]) == (Y_DN, num_nodes)

# Get the output graph edges of the 2nd graph
@assert size(y.gf[:,2]) == (Y_DG,)

Example 2: Batch of graphs with different structures

adj_mat_1 = [
    1 0 1;
    1 1 0;
    0 0 1;
] # Adjacency matrix 1
num_nodes_1 = size(adj_mat_1, 1)
num_edges_1 = length(filter(isone, adj_mat_1))

adj_mat_2 = [
    1 0 1 0;
    1 1 0 1;
    0 0 1 0;
    1 1 0 1;
] # Adjacency matrix 2
num_nodes_2 = size(adj_mat_2, 1)
num_edges_2 = length(filter(isone, adj_mat_2))

edge_features = [
    rand(Float32, X_DE, num_edges_1),
    rand(Float32, X_DE, num_edges_2),
]
node_features = [
    rand(Float32, X_DN, num_nodes_1),
    rand(Float32, X_DN, num_nodes_2),
]
graph_features = nothing # no graph level input features

x = (
    graphs=[adj_mat_1,adj_mat_2],  # Graphs in this batch have different structure
    ef=edge_features, 
    nf=node_features,
    gf=graph_features
) |> batch

y_batched = block(x)
y = y_batched |> unbatch

# Memory-efficient view of features for a batch with different graph structures
@assert size(efview(y_batched, :, :, 1)) == (Y_DE, num_edges_1) # edge features for graph 1
@assert size(nfview(y_batched, :, :, 1)) == (Y_DN, num_nodes_1)  # edge features for graph 1
@assert size(gfview(y_batched, :, 1)) == (Y_DG,) # graph features for graph 1
@assert size(efview(y_batched, :, :, 2)) == (Y_DE, num_edges_2) # edge features for graph 2
@assert size(nfview(y_batched, :, :, 2)) == (Y_DN, num_nodes_2) # node features for graph 2
@assert size(gfview(y_batched, :, 2)) == (Y_DG,) # graph features for graph 2

# Copied array of features (less efficient) for a batch with different graph structures
@assert size(y.ef[1]) == (Y_DE, num_edges_1) # edge features for graph 1
@assert size(y.nf[1]) == (Y_DN, num_nodes_1)  # edge features for graph 1
@assert size(y.gf[1]) == (Y_DG,) # graph features for graph 1
@assert size(y.ef[2]) == (Y_DE, num_edges_2) # edge features for graph 2
@assert size(y.nf[2]) == (Y_DN, num_nodes_2) # node features for graph 2
@assert size(y.gf[2]) == (Y_DG,) # graph features for graph 2

Example 3: Encoder -> Core -> Decoder (sequential blocks)

input_dims = (X_DE, X_DN, X_DG)
core_dims = (10, 5, 3)
output_dims = (Y_DE, Y_DN, Y_DG)

struct GNNModel{E,C,D}
    encoder::E
    core_list::C
    decoder::D
end

function GNNModel(; n_cores=2)
    GNNModel(
        GNBlock(input_dims => core_dims),
        GNCoreList([GNCore(core_dims) for _ in 1:n_cores]),
        GNBlock(core_dims => output_dims),
    )
end

function (m::GNNModel)(x)
    (m.decoder ∘ m.core_list ∘ m.encoder)(x)
end

m = GNNModel()

adj_mat = [
    1 0 1;
    1 1 0;
    0 0 1;
]

num_nodes = size(adj_mat, 1)
num_edges = length(filter(isone, adj_mat))

batch_size = 2
edge_features = rand(Float32, X_DE, num_edges, batch_size)
node_features = rand(Float32, X_DN, num_nodes, batch_size)
graph_features = nothing # no graph level input features

x = (
    graphs=adj_mat, # All graphs in this batch have same structure
    ef=edge_features, # (X_DE, num_edges, batch_size)
    nf=node_features, # (X_DN, num_nodes, batch_size)
    gf=graph_features # (X_DG, batch_size)
) |> batch

y = block(x) |> unbatch

@assert size(y.ef) == (Y_DE, num_edges, batch_size)
@assert size(y.nf) == (Y_DN, num_nodes, batch_size)
@assert size(y.gf) == (Y_DG, batch_size)

Components

GraphNets.GNBlock — Type

GNBlock(in => out;  dropout=0)

Initializes an instance of the GNBlock type, representing a GraphNet block.

The following keyword arguments are supported:

dropout (Defaults to 0)

Example:

in_dims = (X_DE, X_DN, X_DG) = 10, 5, 0
out_dims = (Y_DE, Y_DN, Y_DG) = 3, 4, 5
block = GNBlock(in_dims => out_dims)
adj_mat = adj_mat = [
    1 0 1;
    1 1 0;
    0 0 1;
]
num_nodes = size(adj_mat, 1)
num_edges = length(filter(isone, adj_mat))
batch_size = 2
edge_features = rand(Float32, X_DE, num_edges, batch_size)
node_features = rand(Float32, X_DN, num_nodes, batch_size)
graph_features = nothing # no graph level input features
x = (
    graphs=adj_mat, # All graphs in this batch have same structure
    ef=edge_features, # (X_DE, num_edges, batch_size)
    nf=node_features, # (X_DN, num_nodes, batch_size)
    gf=graph_features # no input graph features
) |> batch
y = block(x) |> unbatch
@assert size(y.ef) == (Y_DE, num_edges, batch_size)
@assert size(y.nf) == (Y_DN, num_nodes, batch_size)
@assert size(y.gf) == (Y_DG, batch_size)

source

GraphNets.GNCore — Type

GNCore(dims; dropout=0)

Initializes an instance of the GNCore type, representing a GraphNet "core" block.

The following keyword arguments are supported:

dropout (Defaults to 0)

Example:

dims = (DE, DN, DG) = 3, 4, 5
core = GNCore(dims)
adj_mat = adj_mat = [
    1 0 1;
    1 1 0;
    0 0 1;
]
num_nodes = size(adj_mat, 1)
num_edges = length(filter(isone, adj_mat))
batch_size = 2
edge_features = rand(Float32, DE, num_edges, batch_size)
node_features = rand(Float32, DN, num_nodes, batch_size)
graph_features = rand(Float32, DG, batch_size)
x = (
    graphs=adj_mat, # All graphs in this batch have same structure
    ef=edge_features, # (DE, num_edges, batch_size)
    nf=node_features, # (DN, num_nodes, batch_size)
    gf=graph_features # no input graph features
) |> batch
y = core(x) |> unbatch
@assert size(y.ef) == (DE, num_edges, batch_size)
@assert size(y.nf) == (DN, num_nodes, batch_size)
@assert size(y.gf) == (DG, batch_size)

source

GraphNets.GNCoreList — Type

GNCoreList(input_dim, num_heads; head_size=(input_dim ÷ num_heads), dropout=0)

Initializes an instance of the GNCoreList type, representing a sequence of GraphNet core blocks composed together.

The following keyword arguments are supported:

dropout (Defaults to 0)

Example:

dims = (DE, DN, DG) = 3, 4, 5
core_list = GNCoreList([GNCore(dims), GNCore(dims)])
adj_mat = adj_mat = [
    1 0 1;
    1 1 0;
    0 0 1;
]
num_nodes = size(adj_mat, 1)
num_edges = length(filter(isone, adj_mat))
batch_size = 2
edge_features = rand(Float32, DE, num_edges, batch_size)
node_features = rand(Float32, DN, num_nodes, batch_size)
graph_features = rand(Float32, DG, batch_size)
x = (
    graphs=adj_mat, # All graphs in this batch have same structure
    ef=edge_features, # (DE, num_edges, batch_size)
    nf=node_features, # (DN, num_nodes, batch_size)
    gf=graph_features # no input graph features
) |> batch
y = core_list(x) |> unbatch
@assert size(y.ef) == (DE, num_edges, batch_size)
@assert size(y.nf) == (DN, num_nodes, batch_size)
@assert size(y.gf) == (DG, batch_size)

source

GraphNets.batch — Function

batch(t::NamedTuple)

Example:

dims = (DE, DN, DG) = 3, 4, 5
core = GNCore(dims)
adj_mat_1 = [
    1 0 1;
    1 1 0;
    0 0 1;
] # Adjacency matrix 1
num_nodes_1 = size(adj_mat_1, 1)
num_edges_1 = length(filter(isone, adj_mat_1))
adj_mat_2 = [
    1 0 1 0;
    1 1 0 1;
    0 0 1 0;
    1 1 0 1;
] # Adjacency matrix 2
num_nodes_2 = size(adj_mat_2, 1)
num_edges_2 = length(filter(isone, adj_mat_2))
edge_features = [
    rand(Float32, DE, num_edges_1),
    rand(Float32, DE, num_edges_2),
]
node_features = [
    rand(Float32, DN, num_nodes_1),
    rand(Float32, DN, num_nodes_2),
]
graph_features = [
    rand(Float32, DG),
    rand(Float32, DG),
]
adj_mats = [adj_mat_1,adj_mat_2]
batch_size = length(adj_mats)
x = (
    graphs=adj_mats,  # Graphs in this batch have different structure
    ef=edge_features, 
    nf=node_features,
    gf=graph_features
)
x_batched = batch(x)
edge_block_size = x_batched.graphs.edge_block_size
node_block_size = x_batched.graphs.node_block_size
@assert typeof(x_batched.graphs) == GraphNets.GNGraphBatch
@assert size(x_batched.ef) == (DE, edge_block_size, batch_size)
@assert size(x_batched.nf) == (DN, node_block_size, batch_size)
@assert size(x_batched.gf) == (DG, 1, batch_size)

source

GraphNets.unbatch — Function

unbatch(t::NamedTuple)

unbatch is the inverse of batch. It takes a batched tuple and returns a tuple of unbatched arrays.

source

GraphNets.efview — Function

efview(t::NamedTuple, d1, d2, d3)

Returns an array view of the edge features contained in the batched output.

source

GraphNets.nfview — Function

nfview(t::NamedTuple, d1, d2, d3)

Returns an array view of the node features contained in the batched output.

source

GraphNets.gfview — Function

gfview(t::NamedTuple, d1, d2)

Returns an array view of the graph features contained in the batched output.

source