Function fitting with PyTorch

The purpose of this tutorial is to explain how to embed a neural network model from PyTorch into JuMP.

Required packages

This tutorial requires the following packages

using JuMP
using Test
import Ipopt
import MathOptAI
import Plots
import PythonCall

Training a model

The following script builds and trains a simple neural network in PyTorch. For simplicity, we do not evaluate out-of-sample test performance, or use a batched data loader. In general, you should train your model in Python, and then use torch.save(model, filename) to save it to a .pt file for later use in Julia.

The model is unimportant, but for this example, we are trying to fit noisy observations of the function $f(x) = x^2 - 2x$.

In Python, we ran:

#!/usr/bin/python3
import torch
model = torch.nn.Sequential(
    torch.nn.Linear(1, 16),
    torch.nn.ReLU(),
    torch.nn.Linear(16, 1),
)

n = 1024
x = torch.arange(-2, 2 + 4 / (n - 1), 4 / (n - 1)).reshape(n, 1)
loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
for epoch in range(100):
    optimizer.zero_grad()
    N = torch.normal(torch.zeros(n, 1), torch.ones(n, 1))
    y = x ** 2 -2 * x + 0.1 * N
    loss = loss_fn(model(x), y)
    loss.backward()
    optimizer.step()

torch.save(model, "model.pt")

JuMP model

Our goal for this JuMP model is to load the Neural Network from PyTorch into the objective function, and then minimize the objective for different fixed values of x to recreate the function that the Neural Network has learned to approximate.

First, create a JuMP model:

model = Model(Ipopt.Optimizer)
set_silent(model)
@variable(model, x)

\[ x \]

Then, load the model from PyTorch using MathOptAI.PytorchModel:

predictor = MathOptAI.PytorchModel(joinpath(@__DIR__, "model.pt"))
y, _ = MathOptAI.add_predictor(model, predictor, [x])
@objective(model, Min, only(y))

\[ moai\_Affine_{1} \]

Now, visualize the fitted function y = predictor(x) by repeatedly solving the optimization problem for different fixed values of x:

X, Y = -2:0.1:2, Float64[]
@constraint(model, c, x == 0.0)
for xi in X
    set_normalized_rhs(c, xi)
    optimize!(model)
    @test is_solved_and_feasible(model)
    push!(Y, objective_value(model))
end
Plots.plot(x -> x^2 - 2x, X; label = "Truth", linestype = :dot)
Plots.plot!(X, Y; label = "Fitted")

This page was generated using Literate.jl.