Function fitting with PyTorch

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

Python integration

MathOptAI uses PythonCall.jl to call from Julia into Python.

To use PytorchModel your code must load the PythonCall package:

import PythonCall

PythonCall uses CondaPkg.jl to manage Python dependencies. See CondaPkg.jl for more control over how to link Julia to an existing Python environment. For example, if you have an existing Python installation (with PyTorch installed), and it is available in the current Conda environment, do:

ENV["JULIA_CONDAPKG_BACKEND"] = "Current"
import PythonCall

If the Python installation can be found on the path and it is not in a Conda environment, do:

ENV["JULIA_CONDAPKG_BACKEND"] = "Null"
import PythonCall

If python is not on your path, you may additionally need to set JULIA_PYTHONCALL_EXE, for example, do:

ENV["JULIA_PYTHONCALL_EXE"] = "python3"
ENV["JULIA_CONDAPKG_BACKEND"] = "Null"
import PythonCall

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.