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.AvgPool2d
• nn.Conv2d
• nn.Flatten
• nn.GELU
• nn.LayerNorm
• nn.LeakyReLU
• nn.Linear
• nn.MaxPool2d
• nn.ReLU
• nn.Sigmoid
• nn.Softmax
• 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")

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[1] - max(0, moai_Affine[1]) = 0
  ├ moai_ReLU[2] ≥ 0
  └ moai_ReLU[2] - max(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.4644489884376526 * max(0, 0.7473132610321045 x[1] - 0.27370238304138184))) + (0.5382549166679382 * max(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 vector 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_Pytorch[1]
julia> formulationMathOptAI.GrayBox{MathOptAI.PytorchModel}(MathOptAI.PytorchModel("saved_pytorch_model.pt"), "cpu", true)
├ variables [1]
│ └ moai_Pytorch[1]
└ constraints [1]
  └ [x[1], moai_Pytorch[1]] ∈ VectorNonlinearOracle{Float64}(;
    dimension = 2,
    l = [0.0],
    u = [0.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 [8]
  ├ moai_ReLU[1] ≥ 0
  ├ moai_z[1] ≥ 0
  ├ moai_Affine[1] - moai_ReLU[1] + moai_z[1] = 0
  ├ [moai_ReLU[1], moai_z[1]] ∈ MathOptInterface.SOS1{Float64}([1.0, 2.0])
  ├ moai_ReLU[2] ≥ 0
  ├ moai_z[2] ≥ 0
  ├ moai_Affine[2] - moai_ReLU[2] + moai_z[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

Custom layers

If your PyTorch model contains a custom layer, define a new AbstractPredictor and pass a config dictionary that maps the Class object to a callback that builds the new predictor.

The callback must have the signature (layer::PythonCall.Py; kwargs...). Valid keyword arguments are currently:

• input_size: the input size of they layer
• config: the config dictionary, if needed to convert layers inside the custom layer
• nn: a reference to torch.nn

You must always have kwargs... so that future versions of MathOptAI can add new keywords in a non-breaking way.

julia> using JuMP, PythonCall, MathOptAI
julia> dir = mktempdir()"/tmp/jl_O65XF9"
julia> write(
           joinpath(dir, "custom_model.py"),
           """
           import torch
           class Skip(torch.nn.Module):
               def __init__(self, inner):
                   super().__init__()
                   self.inner = inner
               def forward(self, x):
                   return self.inner(x) + x
           """,
       )186
julia> filename = joinpath(dir, "custom_model.pt")"/tmp/jl_O65XF9/custom_model.pt"
julia> PythonCall.@pyexec(
           (dir, filename) =>
               """
               import sys
               sys.path.insert(0, dir)
               import torch
               from custom_model import Skip
               inner = torch.nn.Sequential(torch.nn.Linear(3, 3), torch.nn.ReLU())
               model = Skip(inner)
               torch.save(model, filename)
               """ => Skip,
       )Python: <class 'custom_model.Skip'>
julia> struct CustomPredictor <: MathOptAI.AbstractPredictor
           p::MathOptAI.Pipeline
       end
julia> function MathOptAI.add_predictor(
           model::JuMP.AbstractModel,
           predictor::CustomPredictor,
           x::Vector;
           kwargs...,
       )
           y, formulation = MathOptAI.add_predictor(model, predictor.p, x; kwargs...)
           @assert length(x) == length(y)
           return y .+ x, formulation
       end
julia> model = Model();
julia> @variable(model, x[i in 1:3]);
julia> predictor = MathOptAI.PytorchModel(filename)MathOptAI.PytorchModel("/tmp/jl_O65XF9/custom_model.pt")
julia> function skip_callback(layer::PythonCall.Py; input_size, kwargs...)
           return CustomPredictor(MathOptAI.build_predictor(layer.inner))
       endskip_callback (generic function with 1 method)
julia> config = Dict(Skip => skip_callback)Dict{PythonCall.Py, typeof(Main.skip_callback)} with 1 entry:
  <class 'custom_model.Skip'> => skip_callback
julia> y, formulation = MathOptAI.add_predictor(model, predictor, x; config);
julia> y3-element Vector{JuMP.AffExpr}:
 moai_ReLU[1] + x[1]
 moai_ReLU[2] + x[2]
 moai_ReLU[3] + x[3]
julia> formulationAffine(A, b) [input: 3, output: 3]
├ variables [3]
│ ├ moai_Affine[1]
│ ├ moai_Affine[2]
│ └ moai_Affine[3]
└ constraints [3]
  ├ -0.5295092463493347 x[1] - 0.4741993248462677 x[2] + 0.29495978355407715 x[3] - moai_Affine[1] = 0.052499689161777496
  ├ 0.3080519139766693 x[1] - 0.29857560992240906 x[2] - 0.23944857716560364 x[3] - moai_Affine[2] = -0.4475747346878052
  └ -0.5054284334182739 x[1] - 0.15384532511234283 x[2] + 0.10201524198055267 x[3] - moai_Affine[3] = -0.5544887185096741
MathOptAI.ReLU()
├ variables [3]
│ ├ moai_ReLU[1]
│ ├ moai_ReLU[2]
│ └ moai_ReLU[3]
└ constraints [6]
  ├ moai_ReLU[1] ≥ 0
  ├ moai_ReLU[1] - max(0, moai_Affine[1]) = 0
  ├ moai_ReLU[2] ≥ 0
  ├ moai_ReLU[2] - max(0, moai_Affine[2]) = 0
  ├ moai_ReLU[3] ≥ 0
  └ moai_ReLU[3] - max(0, moai_Affine[3]) = 0