Hi,
I’m the author of GraphNeuralNetworks.jl (GNN.jl). I wonder why rolling out a new and very limited gnn package instead of contributing to existing ones and benefitting the community. If we want to bring the julia ML ecosystem, and specifically the part related to gnns, up to the standards of python alternatives we have to join forces.
It looks like GraphNeuralNetworks.jl offers GNN.jl implementations primary of the convolutional and attentional type. It does seem there is a message-passing variant, but I don’t see that edge vectors are allowed as input or output (I see scalar edge weights, but no support for full edge feature vectors).
Actually many layers in GNN.jl support scalar and vectors edge features, see the table in the docs. Also, GNN.jl allows to easily express the message passing layer you implemented.
GraphNets.jl is an implementation of a particular GNN architecture, Graphs Nets, which is a message passing variety of GNN that allows input/output feature vectors of nodes, edges, and at the graph/aggregate level. As GraphNets.jl is not a tool that supports other GNN flavors (convolutional, attention, and other message passing varieties), specialization can be applied that allows an optimized implementation of only the “Graph Nets” architecture as outlined here .
Did you try to implement the functionality you want using GNN.jl?
From a brief peek at the library and at the paper " Relational inductive biases, …" (which btw is mainly a position paper, the message passing it proposes is not much used in practice nor it is implemented by the main python graph library) your GNNBlock can be implemented in GNN.jl as (code not tested)
using GraphNeuralNetworks
struct GNBlock{E,N,G}
edgefn::E
nodefn::N
graphfn::G
end
@functor GNBlock
function GNBlock(((ein, nin, gin), (eout, nout, gout)))
edge_input_in = edge_in + 2*node_in + graph_in
node_input_in = node_in + edge_out + graph_in
graph_input_in = node_out + edge_out + graph_in
return GNBlock(Dense(edge_input_in, edge_out),
Dense(node_input_in, node_out),
Dense(graph_input_in, graph_out))
end
function (m::GNBlock)(g::GNNGraph, x, e, u)
e = vcat(e, broadcast_edges(g, u))
e = apply_edges((xi, xj, e) -> vcat(xi, xj, e); x, e)
e = m.edgefn(e)
ē = aggregate_neighbors(g, mean, e)
x = m.nodefn(vcat(x, ē, broadcast_nodes(g, u)))
u = vcat(reduce_nodes(mean, g, x), reduce_edges(mean, g, e), u)
u = m.graphfn(u)
return x, e, u
end
This is something worthy of a PR to GNN.jl. If you felt some features were missing, it would have been good to open an issue, start a discussion …
Now we would have a library with 19 layers, instead of 2 libraries, one with 18 layers and one with 1 layer.
In particular, I think “GraphNets.jl” has a novel approach of batching together graphs of different structure (different adjacency matrices), that allows optimal training speed of mini-batches and also allows all matrix multiplications to take advantage of cuBLAS Strided Batched Matrix Multiply (which basically means using NNlib.batched_mul).
How does this batching works? Is it documented somewhere? One should benchmark the performance on standard tasks and datasets, e.g. the examples in
From my glimpse at the library, I suspect that the message passing scheme doesn’t scale well since it involves dense matrix multiplications. You cannot go beyond small/medium graphs with this approach.
Efficient message passing approaches are typically implemented in terms of gather/scatter operations or generalized sparse matrix multiplications.
