Graph neural networks

The purpose of this tutorial is to explain how to embed graph neural network models from PyTorch Geometric into JuMP.

Required packages

This tutorial requires the following packages

using JuMP
using Test
import Ipopt
import MathOptAI
import PythonCall

Setting up the Python side

We have written a custom PyTorch module, which is stored in the file my_gnn.py:

print(read(joinpath(@__DIR__, "my_gnn.py"), String))

import torch
import torch_geometric

class MyGNN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = torch_geometric.nn.GCNConv(2, 3)

    def forward(self, x, edge_index):
        x = self.conv(x, edge_index)
        return torch.nn.functional.sigmoid(x)

To make it easy to load this Python file into the documentation, we add the current directory to the Python path:

dir = @__DIR__
PythonCall.@pyexec(dir => "import sys; sys.path.insert(0, dir)")

Then, we can import torch and the my_gnn file into Julia:

my_gnn = PythonCall.pyimport("my_gnn");

Now, we can load our GNN into Julia:

predictor = my_gnn.MyGNN()

Python:
MyGNN(
  (conv): GCNConv(2, 3)
)

In practice, you should train the GNN in Python, and write it to a file using torch.save; see the Function fitting with PyTorch for details. You could then replace the predictor with an appropriate PytorchModel.

Creating the graph

Before we can embed predictor into a JuMP model, we need to specialize it to a particular graph structure, and we need to provide a callback function that MathOptAI can use to turn an instance of MyGNN into an AbstractPredictor.

First, we need the graph structure. Rather than providing a 2-by-n tensor, we provide the edges as a list of i => j pairs:

edge_index = [1 => 2, 2 => 1, 2 => 3, 3 => 2, 3 => 4, 4 => 3]

6-element Vector{Pair{Int64, Int64}}:
 1 => 2
 2 => 1
 2 => 3
 3 => 2
 3 => 4
 4 => 3

Note that the nodes are 1-indexed.

The callback takes in a PythonCall.Py layer, and returns an AbstractPredictor that matches the forward pass of the GNN:

function MyGNN_callback(layer::PythonCall.Py; kwargs...)
    return MathOptAI.Pipeline(
        MathOptAI.GCNConv(layer.conv; edge_index),
        MathOptAI.Sigmoid(),
    )
end

MyGNN_callback (generic function with 1 method)

Note how the callback uses edge_index.

Creating the JuMP model

Now we can embed predictor into a JuMP model.

model = Model(Ipopt.Optimizer)
set_silent(model)

Our input x has four rows, one for each node, and two columns, one for each input attribute. The fixed bounds are so we can easily compare solutions between Python and Julia; replace them with your real constraints in practice.

@variable(model, x[i in 1:4, j in 1:2] == sin(i) + cos(j))

4×2 Matrix{JuMP.VariableRef}:
 x[1,1]  x[1,2]
 x[2,1]  x[2,2]
 x[3,1]  x[3,2]
 x[4,1]  x[4,2]

Then, we add predictor with the config argument mapping gnn.MyGNN (the Python class) to our callback MyGNN_callback:

y, _ = MathOptAI.add_predictor(
    model,
    predictor,
    x;
    config = Dict(my_gnn.MyGNN => MyGNN_callback),
);

Now we can solve the model:

optimize!(model)
assert_is_solved_and_feasible(model)

and look at the solution of y:

Y = reshape(value(y), 4, 3)

4×3 Matrix{Float64}:
 0.540946  0.560852  0.649995
 0.559764  0.586849  0.65636
 0.589818  0.626102  0.603036
 0.592409  0.628382  0.556882

Which is identical to what we obtain when we evaluate the GNN in Python:

torch = PythonCall.pyimport("torch")
py_x = torch.tensor([value.(x[i, :]) for i in 1:4])
# `py_edge_index` is the same graph as `edge_index`, just in a different form.
# Note also that this is 0-indexed.
py_edge_index = torch.tensor([[0, 1, 1, 2, 2, 3], [1, 0, 2, 1, 3, 2]])
predictor(py_x, py_edge_index)

Python:
tensor([[0.5409, 0.5609, 0.6500],
        [0.5598, 0.5868, 0.6564],
        [0.5898, 0.6261, 0.6030],
        [0.5924, 0.6284, 0.5569]], grad_fn=<SigmoidBackward0>)

This page was generated using Literate.jl.