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, PythonCalljulia> 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.08490216732025146MathOptAI.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]) = 0Affine(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, PythonCalljulia> 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.4644489884376526 * max(0, 0.7473132610321045 x[1] - 0.27370238304138184) + 0.5382549166679382 * max(0, -0.8237636089324951 x[1] - 0.08490216732025146)) + -0.5698285102844238julia> 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, PythonCalljulia> 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> formulationMathOptAI.GrayBox{MathOptAI.PytorchModel}(MathOptAI.PytorchModel("saved_pytorch_model.pt"), "cpu", true)├ variables [1]│ └ moai_GrayBox[1]└ constraints [1]  └ [x[1], moai_GrayBox[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, PythonCalljulia> 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.08490216732025146MathOptAI.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, MathOptAIjulia> dir = mktempdir()"/tmp/jl_JAQQoO"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           """,       )186julia> filename = joinpath(dir, "custom_model.pt")"/tmp/jl_JAQQoO/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       endjulia> 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       endjulia> model = Model();julia> @variable(model, x[i in 1:3]);julia> predictor = MathOptAI.PytorchModel(filename)MathOptAI.PytorchModel("/tmp/jl_JAQQoO/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_callbackjulia> 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.43555766344070435 x[1] + 0.13670268654823303 x[2] - 0.1979544311761856 x[3] - moai_Affine[1] = -0.38053420186042786  ├ 0.5095124840736389 x[1] - 0.31588977575302124 x[2] + 0.5026256442070007 x[3] - moai_Affine[2] = 0.5160393714904785  └ -0.33771297335624695 x[1] - 0.42213013768196106 x[2] - 0.23300319910049438 x[3] - moai_Affine[3] = -0.164980947971344MathOptAI.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