Советы
Пользовательские операторы с векторными выходными данными
Распространенной ситуацией является наличие пользовательского оператора наподобие следующего, который возвращает несколько элементов выходных данных (для подсчета того, сколько раз вызывается этот метод, мы определяем function_calls
):
function_calls = 0
function foo(x, y)
global function_calls += 1
common_term = x^2 + y^2
term_1 = sqrt(1 + common_term)
term_2 = common_term
return term_1, term_2
end
foo (generic function with 1 method)
Например, первый член может использоваться в целевой функции, а второй — в ограничении, и часто их обработка сопряжена с большими вычислительными затратами.
Это проблема для JuMP, так как пользовательские операторы должны возвращать одно число. Одним из вариантов является определение двух отдельных функций, первая из которых возвращает первый аргумент, а вторая — второй.
foo_1(x, y) = foo(x, y)[1]
foo_2(x, y) = foo(x, y)[2]
foo_2 (generic function with 1 method)
Однако если вычисление общего члена требует больших затрат, такой подход является расточительным, так как проблемный член вычисляется дважды. Давайте посмотрим, сколько раз вычисляется x^2 + y^2
во время решения:
model = Model(Ipopt.Optimizer)
set_silent(model)
@variable(model, x[1:2] >= 0, start = 0.1)
@operator(model, op_foo_1, 2, foo_1)
@operator(model, op_foo_2, 2, foo_2)
@objective(model, Max, op_foo_1(x[1], x[2]))
@constraint(model, op_foo_2(x[1], x[2]) <= 2)
function_calls = 0
optimize!(model)
@assert is_solved_and_feasible(model)
Test.@test objective_value(model) ≈ √3 atol = 1e-4
Test.@test value.(x) ≈ [1.0, 1.0] atol = 1e-4
println("Naive approach: function calls = $(function_calls)")
Naive approach: function calls = 44
Альтернативный подход заключается в использовании запоминания, то есть сохранения результатов вычисления функции в кэше. Функцию запоминания можно реализовать так:
"""
memoize(foo::Function, n_outputs::Int)
Take a function `foo` and return a vector of length `n_outputs`, where element
`i` is a function that returns the equivalent of `foo(x...)[i]`.
To avoid duplication of work, cache the most-recent evaluations of `foo`.
Because `foo_i` is auto-differentiated with ForwardDiff, our cache needs to
work when `x` is a `Float64` and a `ForwardDiff.Dual`.
"""
function memoize(foo::Function, n_outputs::Int)
last_x, last_f = nothing, nothing
last_dx, last_dfdx = nothing, nothing
function foo_i(i, x::T...) where {T<:Real}
if T == Float64
if x !== last_x
last_x, last_f = x, foo(x...)
end
return last_f[i]::T
else
if x !== last_dx
last_dx, last_dfdx = x, foo(x...)
end
return last_dfdx[i]::T
end
end
return [(x...) -> foo_i(i, x...) for i in 1:n_outputs]
end
memoize (generic function with 1 method)
Посмотрим, как она работает. Сначала создаются версии foo
с запоминанием:
memoized_foo = memoize(foo, 2)
2-element Vector{Main.var"#4#7"{Int64, Main.var"#foo_i#5"{typeof(Main.foo)}}}:
#4 (generic function with 1 method)
#4 (generic function with 1 method)
Далее производится попытка вычислить первый элемент memoized_foo
.
function_calls = 0
memoized_foo[1](1.0, 1.0)
println("function_calls = ", function_calls)
function_calls = 1
Как и ожидалось, функция была вычислена один раз. Однако если снова вызвать функцию, то произойдет успешное обращение к кэшу вместо повторного вычисления foo
и function_calls
по-прежнему будет равно 1
!
memoized_foo[1](1.0, 1.0)
println("function_calls = ", function_calls)
function_calls = 1
Теперь посмотрим, как это работает во время реального решения:
model = Model(Ipopt.Optimizer)
set_silent(model)
@variable(model, x[1:2] >= 0, start = 0.1)
@operator(model, op_foo_1, 2, memoized_foo[1])
@operator(model, op_foo_2, 2, memoized_foo[2])
@objective(model, Max, op_foo_1(x[1], x[2]))
@constraint(model, op_foo_2(x[1], x[2]) <= 2)
function_calls = 0
optimize!(model)
@assert is_solved_and_feasible(model)
Test.@test objective_value(model) ≈ √3 atol = 1e-4
Test.@test value.(x) ≈ [1.0, 1.0] atol = 1e-4
println("Memoized approach: function_calls = $(function_calls)")
Memoized approach: function_calls = 22
По сравнению с примитивным подходом подход с запоминанием требует вдвое меньше вычислений функции.