Skip to content

Commit 758fdf7

Browse files
odowodow
authored
Add build_predictor (#78)
1 parent fea7a92 commit 758fdf7

File tree

8 files changed

+302
-30
lines changed

8 files changed

+302
-30
lines changed

‎docs/src/api.md

Lines changed: 17 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -31,16 +31,9 @@ AbstractPredictor
3131
add_predictor
3232
```
3333

34-
```@autodocs
35-
Modules = [
36-
Base.get_extension(MathOptAI, :MathOptAIAbstractGPsExt),
37-
Base.get_extension(MathOptAI, :MathOptAIDecisionTreeExt),
38-
Base.get_extension(MathOptAI, :MathOptAIFluxExt),
39-
Base.get_extension(MathOptAI, :MathOptAIGLMExt),
40-
Base.get_extension(MathOptAI, :MathOptAILuxExt),
41-
Base.get_extension(MathOptAI, :MathOptAIPythonCallExt),
42-
Base.get_extension(MathOptAI, :MathOptAIStatsModelsExt),
43-
]
34+
## `build_predictor`
35+
```@docs
36+
build_predictor
4437
```
4538

4639
## `Affine`
@@ -112,3 +105,17 @@ SoftPlus
112105
```@docs
113106
Tanh
114107
```
108+
109+
## Extensions
110+
111+
```@autodocs
112+
Modules = [
113+
Base.get_extension(MathOptAI, :MathOptAIAbstractGPsExt),
114+
Base.get_extension(MathOptAI, :MathOptAIDecisionTreeExt),
115+
Base.get_extension(MathOptAI, :MathOptAIFluxExt),
116+
Base.get_extension(MathOptAI, :MathOptAIGLMExt),
117+
Base.get_extension(MathOptAI, :MathOptAILuxExt),
118+
Base.get_extension(MathOptAI, :MathOptAIPythonCallExt),
119+
Base.get_extension(MathOptAI, :MathOptAIStatsModelsExt),
120+
]
121+
```

‎ext/MathOptAIDecisionTreeExt.jl

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,10 +53,44 @@ function MathOptAI.add_predictor(
5353
predictor::DecisionTree.Root,
5454
x::Vector,
5555
)
56-
inner_predictor = _tree_or_leaf(predictor.node)
56+
inner_predictor = MathOptAI.build_predictor(predictor.node)
5757
return MathOptAI.add_predictor(model, inner_predictor, x)
5858
end
5959

60+
"""
61+
MathOptAI.build_predictor(predictor::DecisionTree.Root)
62+
63+
Convert a binary decision tree from DecisionTree.jl to a
64+
[`BinaryDecisionTree`](@ref).
65+
66+
## Example
67+
68+
```jldoctest
69+
julia> using MathOptAI, DecisionTree
70+
71+
julia> truth(x::Vector) = x[1] <= 0.5 ? -2 : (x[2] <= 0.3 ? 3 : 4)
72+
truth (generic function with 1 method)
73+
74+
julia> features = abs.(sin.((1:10) .* (3:4)'));
75+
76+
julia> size(features)
77+
(10, 2)
78+
79+
julia> labels = truth.(Vector.(eachrow(features)));
80+
81+
julia> ml_model = DecisionTree.build_tree(labels, features)
82+
Decision Tree
83+
Leaves: 3
84+
Depth: 2
85+
86+
julia> MathOptAI.build_predictor(ml_model)
87+
BinaryDecisionTree{Float64,Int64} [leaves=3, depth=2]
88+
```
89+
"""
90+
function MathOptAI.build_predictor(predictor::DecisionTree.Root)
91+
return _tree_or_leaf(predictor.node)
92+
end
93+
6094
function _tree_or_leaf(node::DecisionTree.Node{K,V}) where {K,V}
6195
return MathOptAI.BinaryDecisionTree{K,V}(
6296
node.featid,

‎ext/MathOptAIFluxExt.jl

Lines changed: 64 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -67,16 +67,76 @@ function MathOptAI.add_predictor(
6767
config::Dict = Dict{Any,Any}(),
6868
reduced_space::Bool = false,
6969
)
70-
inner_predictor = MathOptAI.Pipeline(MathOptAI.AbstractPredictor[])
71-
for layer in predictor.layers
72-
_add_predictor(inner_predictor, layer, config)
73-
end
70+
inner_predictor = MathOptAI.build_predictor(predictor; config)
7471
if reduced_space
7572
inner_predictor = MathOptAI.ReducedSpace(inner_predictor)
7673
end
7774
return MathOptAI.add_predictor(model, inner_predictor, x)
7875
end
7976

77+
"""
78+
MathOptAI.build_predictor(
79+
predictor::Flux.Chain;
80+
config::Dict = Dict{Any,Any}(),
81+
)
82+
83+
Convert a trained neural network from Flux.jl to a [`Pipeline`](@ref).
84+
85+
## Supported layers
86+
87+
* `Flux.Dense`
88+
* `Flux.softmax`
89+
90+
## Supported activation functions
91+
92+
* `Flux.relu`
93+
* `Flux.sigmoid`
94+
* `Flux.softplus`
95+
* `Flux.tanh`
96+
97+
## Keyword arguments
98+
99+
* `config`: a dictionary that maps `Flux` activation functions to an
100+
[`AbstractPredictor`](@ref) to control how the activation functions are
101+
reformulated.
102+
103+
## Example
104+
105+
```jldoctest
106+
julia> using Flux, MathOptAI
107+
108+
julia> chain = Flux.Chain(Flux.Dense(1 => 16, Flux.relu), Flux.Dense(16 => 1));
109+
110+
julia> MathOptAI.build_predictor(
111+
chain;
112+
config = Dict(Flux.relu => MathOptAI.ReLU()),
113+
)
114+
Pipeline with layers:
115+
* Affine(A, b) [input: 1, output: 16]
116+
* ReLU()
117+
* Affine(A, b) [input: 16, output: 1]
118+
119+
julia> MathOptAI.build_predictor(
120+
chain;
121+
config = Dict(Flux.relu => MathOptAI.ReLUQuadratic()),
122+
)
123+
Pipeline with layers:
124+
* Affine(A, b) [input: 1, output: 16]
125+
* ReLUQuadratic()
126+
* Affine(A, b) [input: 16, output: 1]
127+
```
128+
"""
129+
function MathOptAI.build_predictor(
130+
predictor::Flux.Chain;
131+
config::Dict = Dict{Any,Any}(),
132+
)
133+
inner_predictor = MathOptAI.Pipeline(MathOptAI.AbstractPredictor[])
134+
for layer in predictor.layers
135+
_add_predictor(inner_predictor, layer, config)
136+
end
137+
return inner_predictor
138+
end
139+
80140
_default(::typeof(identity)) = nothing
81141
_default(::Any) = missing
82142
_default(::typeof(Flux.relu)) = MathOptAI.ReLU()

‎ext/MathOptAIGLMExt.jl

Lines changed: 63 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -44,13 +44,35 @@ function MathOptAI.add_predictor(
4444
x::Vector;
4545
reduced_space::Bool = false,
4646
)
47-
inner_predictor = MathOptAI.Affine(GLM.coef(predictor))
47+
inner_predictor = MathOptAI.build_predictor(predictor)
4848
if reduced_space
4949
inner_predictor = MathOptAI.ReducedSpace(inner_predictor)
5050
end
5151
return MathOptAI.add_predictor(model, inner_predictor, x)
5252
end
5353

54+
"""
55+
MathOptAI.build_predictor(predictor::GLM.LinearModel)
56+
57+
Convert a trained linear model from GLM.jl to an [`Affine`](@ref) layer.
58+
59+
## Example
60+
61+
```jldoctest
62+
julia> using GLM, MathOptAI
63+
64+
julia> X, Y = rand(10, 2), rand(10);
65+
66+
julia> model_glm = GLM.lm(X, Y);
67+
68+
julia> MathOptAI.build_predictor(model_glm)
69+
Affine(A, b) [input: 2, output: 1]
70+
```
71+
"""
72+
function MathOptAI.build_predictor(predictor::GLM.LinearModel)
73+
return MathOptAI.Affine(GLM.coef(predictor))
74+
end
75+
5476
"""
5577
MathOptAI.add_predictor(
5678
model::JuMP.Model,
@@ -100,12 +122,50 @@ function MathOptAI.add_predictor(
100122
sigmoid::MathOptAI.AbstractPredictor = MathOptAI.Sigmoid(),
101123
reduced_space::Bool = false,
102124
)
103-
affine = MathOptAI.Affine(GLM.coef(predictor))
104-
inner_predictor = MathOptAI.Pipeline(affine, sigmoid)
125+
inner_predictor = MathOptAI.build_predictor(predictor; sigmoid)
105126
if reduced_space
106127
inner_predictor = MathOptAI.ReducedSpace(inner_predictor)
107128
end
108129
return MathOptAI.add_predictor(model, inner_predictor, x)
109130
end
110131

132+
"""
133+
MathOptAI.build_predictor(
134+
predictor::GLM.GeneralizedLinearModel{
135+
GLM.GlmResp{Vector{Float64},GLM.Bernoulli{Float64},GLM.LogitLink},
136+
};
137+
sigmoid::MathOptAI.AbstractPredictor = MathOptAI.Sigmoid(),
138+
)
139+
140+
Convert a trained logistic model from GLM.jl to a [`Pipeline`](@ref) layer.
141+
142+
## Keyword arguments
143+
144+
* `sigmoid`: the predictor to use for the sigmoid layer.
145+
146+
## Example
147+
148+
```jldoctest
149+
julia> using GLM, MathOptAI
150+
151+
julia> X, Y = rand(10, 2), rand(Bool, 10);
152+
153+
julia> model_glm = GLM.glm(X, Y, GLM.Bernoulli());
154+
155+
julia> MathOptAI.build_predictor(model_glm)
156+
Pipeline with layers:
157+
* Affine(A, b) [input: 2, output: 1]
158+
* Sigmoid()
159+
```
160+
"""
161+
function MathOptAI.build_predictor(
162+
predictor::GLM.GeneralizedLinearModel{
163+
GLM.GlmResp{Vector{Float64},GLM.Bernoulli{Float64},GLM.LogitLink},
164+
};
165+
sigmoid::MathOptAI.AbstractPredictor = MathOptAI.Sigmoid(),
166+
)
167+
affine = MathOptAI.Affine(GLM.coef(predictor))
168+
return MathOptAI.Pipeline(affine, sigmoid)
169+
end
170+
111171
end # module

‎ext/MathOptAILuxExt.jl

Lines changed: 72 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -76,16 +76,84 @@ function MathOptAI.add_predictor(
7676
x::Vector;
7777
config::Dict = Dict{Any,Any}(),
7878
reduced_space::Bool = false,
79+
)
80+
inner_predictor = MathOptAI.build_predictor(predictor; config)
81+
if reduced_space
82+
inner_predictor = MathOptAI.ReducedSpace(inner_predictor)
83+
end
84+
return MathOptAI.add_predictor(model, inner_predictor, x)
85+
end
86+
87+
"""
88+
MathOptAI.build_predictor(
89+
predictor::Tuple{<:Lux.Chain,<:NamedTuple,<:NamedTuple};
90+
config::Dict = Dict{Any,Any}(),
91+
)
92+
93+
Convert a trained neural network from Lux.jl to a [`Pipeline`](@ref).
94+
95+
## Supported layers
96+
97+
* `Lux.Dense`
98+
99+
## Supported activation functions
100+
101+
* `Lux.relu`
102+
* `Lux.sigmoid`
103+
* `Lux.softplus`
104+
* `Lux.tanh`
105+
106+
## Keyword arguments
107+
108+
* `config`: a dictionary that maps `Lux` activation functions to an
109+
[`AbstractPredictor`](@ref) to control how the activation functions are
110+
reformulated.
111+
112+
## Example
113+
114+
```jldoctest; filter=r"[┌|└].+"
115+
julia> using Lux, MathOptAI, Random
116+
117+
julia> rng = Random.MersenneTwister();
118+
119+
julia> chain = Lux.Chain(Lux.Dense(1 => 16, Lux.relu), Lux.Dense(16 => 1))
120+
Chain(
121+
layer_1 = Dense(1 => 16, relu), # 32 parameters
122+
layer_2 = Dense(16 => 1), # 17 parameters
123+
) # Total: 49 parameters,
124+
# plus 0 states.
125+
126+
julia> parameters, state = Lux.setup(rng, chain);
127+
128+
julia> predictor = MathOptAI.build_predictor(
129+
(chain, parameters, state);
130+
config = Dict(Lux.relu => MathOptAI.ReLU()),
131+
)
132+
Pipeline with layers:
133+
* Affine(A, b) [input: 1, output: 16]
134+
* ReLU()
135+
* Affine(A, b) [input: 16, output: 1]
136+
137+
julia> MathOptAI.build_predictor(
138+
(chain, parameters, state);
139+
config = Dict(Lux.relu => MathOptAI.ReLUQuadratic()),
140+
)
141+
Pipeline with layers:
142+
* Affine(A, b) [input: 1, output: 16]
143+
* ReLUQuadratic()
144+
* Affine(A, b) [input: 16, output: 1]
145+
```
146+
"""
147+
function MathOptAI.build_predictor(
148+
predictor::Tuple{<:Lux.Chain,<:NamedTuple,<:NamedTuple};
149+
config::Dict = Dict{Any,Any}(),
79150
)
80151
chain, parameters, _ = predictor
81152
inner_predictor = MathOptAI.Pipeline(MathOptAI.AbstractPredictor[])
82153
for (layer, parameter) in zip(chain.layers, parameters)
83154
_add_predictor(inner_predictor, layer, parameter, config)
84155
end
85-
if reduced_space
86-
inner_predictor = MathOptAI.ReducedSpace(inner_predictor)
87-
end
88-
return MathOptAI.add_predictor(model, inner_predictor, x)
156+
return inner_predictor
89157
end
90158

91159
_default(::typeof(identity)) = nothing

‎ext/MathOptAIPythonCallExt.jl

Lines changed: 33 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -41,10 +41,7 @@ function MathOptAI.add_predictor(
4141
config::Dict = Dict{Any,Any}(),
4242
reduced_space::Bool = false,
4343
)
44-
torch = PythonCall.pyimport("torch")
45-
nn = PythonCall.pyimport("torch.nn")
46-
torch_model = torch.load(predictor.filename)
47-
inner_predictor = _predictor(nn, torch_model, config)
44+
inner_predictor = MathOptAI.build_predictor(predictor; config)
4845
if reduced_space
4946
# If config maps to a ReducedSpace predictor, we'll get a MethodError
5047
# when trying to add the nested redcued space predictors.
@@ -54,6 +51,38 @@ function MathOptAI.add_predictor(
5451
return MathOptAI.add_predictor(model, inner_predictor, x)
5552
end
5653

54+
"""
55+
MathOptAI.build_predictor(
56+
predictor::MathOptAI.PytorchModel;
57+
config::Dict = Dict{Any,Any}(),
58+
)
59+
60+
Convert a trained neural network from Pytorch via PythonCall.jl to a
61+
[`Pipeline`](@ref).
62+
63+
## Supported layers
64+
65+
* `nn.Linear`
66+
* `nn.ReLU`
67+
* `nn.Sequential`
68+
* `nn.Sigmoid`
69+
* `nn.Tanh`
70+
71+
## Keyword arguments
72+
73+
* `config`: a dictionary that maps symbols to an [`AbstractPredictor`](@ref)
74+
to control how the activation functions are reformulated.
75+
"""
76+
function MathOptAI.build_predictor(
77+
predictor::MathOptAI.PytorchModel;
78+
config::Dict = Dict{Any,Any}(),
79+
)
80+
torch = PythonCall.pyimport("torch")
81+
nn = PythonCall.pyimport("torch.nn")
82+
torch_model = torch.load(predictor.filename)
83+
return _predictor(nn, torch_model, config)
84+
end
85+
5786
function _predictor(nn, layer, config)
5887
if Bool(PythonCall.pybuiltins.isinstance(layer, nn.Linear))
5988
weight = mapreduce(vcat, layer.weight.tolist()) do w

0 commit comments

Comments
 (0)