PyTorch

PyTorch is a library for machine learning in Python.

The upstream documentation is available at https://pytorch.org/docs/stable/.

Supported layers

MathOptAI supports embedding a PyTorch models into JuMP if it is a nn.Sequential composed of:

• nn.Linear
• nn.ReLU
• nn.Sigmoid
• nn.Softplus
• nn.Tanh

File format

Use torch.save to save a trained PyTorch model to a .pt file:

#!/usr/bin/python3
import torch
model = torch.nn.Sequential(
    torch.nn.Linear(1, 2),
    torch.nn.ReLU(),
    torch.nn.Linear(2, 1),
)
torch.save(model, "saved_pytorch_model.pt")

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

Basic example

Use MathOptAI.add_predictor to embed a PyTorch model into a JuMP model:

julia> using JuMP, MathOptAI, PythonCall
julia> model = Model();
julia> @variable(model, x[1:1]);
julia> predictor = MathOptAI.PytorchModel("saved_pytorch_model.pt");
julia> y, formulation = MathOptAI.add_predictor(model, predictor, x);
julia> y1-element Vector{JuMP.VariableRef}:
 moai_Affine[1]
julia> formulationAffine(A, b) [input: 1, output: 2]
├ variables [2]
│ ├ moai_Affine[1]
│ └ moai_Affine[2]
└ constraints [2]
  ├ 0.7473132610321045 x[1] - moai_Affine[1] = 0.27370238304138184
  └ -0.8237636089324951 x[1] - moai_Affine[2] = 0.08490216732025146
MathOptAI.ReLU()
├ variables [2]
│ ├ moai_ReLU[1]
│ └ moai_ReLU[2]
└ constraints [4]
  ├ moai_ReLU[1] ≥ 0
  ├ moai_ReLU[2] ≥ 0
  ├ moai_ReLU[1] - max(0.0, moai_Affine[1]) = 0
  └ moai_ReLU[2] - max(0.0, moai_Affine[2]) = 0
Affine(A, b) [input: 2, output: 1]
├ variables [1]
│ └ moai_Affine[1]
└ constraints [1]
  └ -0.4644489884376526 moai_ReLU[1] + 0.5382549166679382 moai_ReLU[2] - moai_Affine[1] = 0.5698285102844238

Reduced-space

Use the reduced_space = true keyword to formulate a reduced-space model:

julia> using JuMP, MathOptAI, PythonCall
julia> model = Model();
julia> @variable(model, x[1:1]);
julia> predictor = MathOptAI.PytorchModel("saved_pytorch_model.pt");
julia> y, formulation =
           MathOptAI.add_predictor(model, predictor, x; reduced_space = true);
julia> y1-element Vector{JuMP.NonlinearExpr}:
 ((+(0.0) + (-0.4644489884376526 * max(0.0, 0.7473132610321045 x[1] - 0.27370238304138184))) + (0.5382549166679382 * max(0.0, -0.8237636089324951 x[1] - 0.08490216732025146))) + -0.5698285102844238
julia> formulationReducedSpace(Affine(A, b) [input: 1, output: 2])
├ variables [0]
└ constraints [0]
ReducedSpace(MathOptAI.ReLU())
├ variables [0]
└ constraints [0]
ReducedSpace(Affine(A, b) [input: 2, output: 1])
├ variables [0]
└ constraints [0]

Gray-box

Use the gray_box = true keyword to embed the network as a nonlinear operator:

julia> using JuMP, MathOptAI, PythonCall
julia> model = Model();
julia> @variable(model, x[1:1]);
julia> predictor = MathOptAI.PytorchModel("saved_pytorch_model.pt");
julia> y, formulation =
           MathOptAI.add_predictor(model, predictor, x; gray_box = true);
julia> y1-element Vector{JuMP.VariableRef}:
 moai_GrayBox[1]
julia> formulationGrayBox
├ variables [1]
│ └ moai_GrayBox[1]
└ constraints [1]
  └ op_##1242(x[1]) - moai_GrayBox[1] = 0

Change how layers are formulated

Pass a dictionary to the config keyword that maps the Symbol name of each PyTorch layer to a MathOptAI predictor:

julia> using JuMP, MathOptAI, PythonCall
julia> model = Model();
julia> @variable(model, x[1:1]);
julia> predictor = MathOptAI.PytorchModel("saved_pytorch_model.pt");
julia> y, formulation = MathOptAI.add_predictor(
           model,
           predictor,
           x;
           config = Dict(:ReLU => MathOptAI.ReLUSOS1()),
       );
julia> y1-element Vector{JuMP.VariableRef}:
 moai_Affine[1]
julia> formulationAffine(A, b) [input: 1, output: 2]
├ variables [2]
│ ├ moai_Affine[1]
│ └ moai_Affine[2]
└ constraints [2]
  ├ 0.7473132610321045 x[1] - moai_Affine[1] = 0.27370238304138184
  └ -0.8237636089324951 x[1] - moai_Affine[2] = 0.08490216732025146
MathOptAI.ReLUSOS1()
├ variables [4]
│ ├ moai_ReLU[1]
│ ├ moai_ReLU[2]
│ ├ moai_z[1]
│ └ moai_z[2]
└ constraints [6]
  ├ moai_ReLU[1] ≥ 0
  ├ moai_ReLU[2] ≥ 0
  ├ moai_Affine[1] - moai_ReLU[1] + moai_z[1] = 0
  ├ moai_Affine[2] - moai_ReLU[2] + moai_z[2] = 0
  ├ [moai_ReLU[1], moai_z[1]] ∈ MathOptInterface.SOS1{Float64}([1.0, 2.0])
  └ [moai_ReLU[2], moai_z[2]] ∈ MathOptInterface.SOS1{Float64}([1.0, 2.0])
Affine(A, b) [input: 2, output: 1]
├ variables [1]
│ └ moai_Affine[1]
└ constraints [1]
  └ -0.4644489884376526 moai_ReLU[1] + 0.5382549166679382 moai_ReLU[2] - moai_Affine[1] = 0.5698285102844238