Minibatches of graphs in GeometricFlux.jl

bad_at_math
Thu, 21 Oct 2021 18:18:08 +0000
Machine Learning

In the supervised graph problem, you train a neural network that maps entire graph instances to target values. The GeometricFlux.jl package has great functionality for training graph neural networks; however, all of the examples (https://github.com/FluxML/GeometricFlux.jl/tree/master/examples) focus on the problem of making node-level predictions (rather than graph-level predictions).

Is there a preferred syntax for using GeometricFlux.jl models to make predictions over minibatches of graphs? A MWE below gives how I’ve been handling this, but I’m not sure if I’m making some dumb or inefficient design decision. My key design detail is probably the definition of a (f::Chain)(x::AbstractVector{FeaturedGraph}) method.

I compare this to the dgl.batch function in Python’s DGL (Training a GNN for Graph Classification — DGL 0.9 documentation), which explicitly constructs a graph minibatch.

Thanks!

using Flux, GeometricFlux, LightGraphs

# generate a graph with node features
function get_graph()
    mg = SimpleGraph()
    add_vertices!(mg, 3)
    add_edge!(mg, 1, 2)
    add_edge!(mg, 2, 3)
    adj_mat = Float32.(adjacency_matrix(mg))
    FeaturedGraph(adj_mat; nf=ones(Float32, 2, nv(mg)))
end

# a method for making predictions over a minibatch of `FeaturedGraph` datapoints
(f::Chain)(x::AbstractVector{FeaturedGraph}) = Flux.stack(f.(x), 1) |> vec

# define a simple graph neural network
function get_model(input_dim=2, hidden_dim=3, output_dim=1)
    Chain(GraphConv(input_dim => hidden_dim, tanh),  # x_i = tanh( Θ_1 x_i + ∑_{j in N(i)} Θ_2 x_j )
          node_feature,
          x -> sum(x, dims=2),  # sum aggregation
          Dense(hidden_dim, output_dim))
end

model = get_model()
# standard loss function for graph regression
loss(x, y) = Flux.Losses.mse(model(x), y)

# a single data point
graph = get_graph()
target = 0.5f0

# minibatch of data points
graphs = FeaturedGraph[graph, graph, graph]
targets = [target, target, target]

# evaluate model on minibatch
loss(graphs, targets)

3 posts - 2 participants