Автоматическое дифференцирование пользовательских операторов

Цель данного руководства — продемонстрировать применение автоматического дифференцирования к пользовательским операторам.

Это руководство предназначено для опытных пользователей. Вместо создания оператора можно воспользоваться трассировкой функций (Function tracing), а если оператор необходим, можно использовать вариант по умолчанию @operator(model, op_f, N, f) вместо явной передачи градиентов и гессианов (Gradients and Hessians).

В этом руководстве используются следующие пакеты.

using JuMP
import Enzyme
import ForwardDiff
import Ipopt
import Test

В качестве простого примера рассмотрим функцию Розенброка:

f(x::T...) where {T} = (1 - x[1])^2 + 100 * (x[2] - x[1]^2)^2

f (generic function with 1 method)

Вот значение в случайной точке:

x = rand(2)

2-element Vector{Float64}:
 0.8809678954777748
 0.5901641516460916

f(x...)

3.4715474597921627

Аналитическая производная

Если выражения достаточно просты, можно предоставить аналитические функции для градиента и гессиана.

Градиент

Функция Розенброка имеет градиентный вектор:

function analytic_∇f(g::AbstractVector, x...)
    g[1] = 400 * x[1]^3 - 400 * x[1] * x[2] + 2 * x[1] - 2
    g[2] = 200 * (x[2] - x[1]^2)
    return
end

analytic_∇f (generic function with 1 method)

Давайте вычислим ее с тем же вектором x:

analytic_g = zeros(2)
analytic_∇f(analytic_g, x...)
analytic_g

2-element Vector{Float64}:
  65.28490308207542
 -37.18805624328958

Гессиан

Гессианова матрица имеет следующий вид:

function analytic_∇²f(H::AbstractMatrix, x...)
    H[1, 1] = 1200 * x[1]^2 - 400 * x[2] + 2
    # H[1, 2] = -400 * x[1] <-- не требуется, так как гессиан симметричен
    H[2, 1] = -400 * x[1]
    H[2, 2] = 200.0
    return
end

analytic_∇²f (generic function with 1 method)

Обратите внимание: поскольку гессиан симметричен, JuMP требует заполнения только нижнего треугольника.

analytic_H = zeros(2, 2)
analytic_∇²f(analytic_H, x...)
analytic_H

2×2 Matrix{Float64}:
  697.26     0.0
 -352.387  200.0

Пример JuMP

Объединив аналитические функции, получим следующее:

function analytic_rosenbrock()
    model = Model(Ipopt.Optimizer)
    set_silent(model)
    @variable(model, x[1:2])
    @operator(model, op_rosenbrock, 2, f, analytic_∇f, analytic_∇²f)
    @objective(model, Min, op_rosenbrock(x[1], x[2]))
    optimize!(model)
    Test.@test is_solved_and_feasible(model)
    return value.(x)
end

analytic_rosenbrock()

2-element Vector{Float64}:
 0.9999999999999718
 0.9999999999999432

ForwardDiff

Вместо аналитических функций для вычисления производных можно использовать пакет ForwardDiff.jl.

Если градиент или гессиан не указан, JuMP по умолчанию использует ForwardDiff.jl для вычисления производных. В этом разделе приводится рабочий пример для демонстрации внутренних процессов.

Преимущества и недостатки

Главное преимущество пакета ForwardDiff в том, что он прост, надежен и поддерживает различные синтаксические конструкции Julia.

Главный же недостаток в том, что функция f должна принимать аргументы x::Real.... Дополнительные сведения см. в разделе Common mistakes when writing a user-defined operator.

Градиент

Градиент можно вычислить с помощью ForwardDiff.gradient!. Обратите внимание, что ForwardDiff ожидает один аргумент Vector{T}, но мы получаем x как кортеж, поэтому потребуется выполнить y -> f(y...) и collect(x), чтобы получить данные в правильном формате.

function fdiff_∇f(g::AbstractVector{T}, x::Vararg{T,N}) where {T,N}
    ForwardDiff.gradient!(g, y -> f(y...), collect(x))
    return
end

fdiff_∇f (generic function with 1 method)

Давайте проверим, получено ли аналитическое решение:

fdiff_g = zeros(2)
fdiff_∇f(fdiff_g, x...)
Test.@test ≈(analytic_g, fdiff_g)

Test Passed

Гессиан

С гессианом ситуация немного сложнее, но код для его реализации выглядит так:

function fdiff_∇²f(H::AbstractMatrix{T}, x::Vararg{T,N}) where {T,N}
    h = ForwardDiff.hessian(y -> f(y...), collect(x))
    for i in 1:N, j in 1:i
        H[i, j] = h[i, j]
    end
    return
end

fdiff_∇²f (generic function with 1 method)

Давайте проверим, получено ли аналитическое решение:

fdiff_H = zeros(2, 2)
fdiff_∇²f(fdiff_H, x...)
Test.@test ≈(analytic_H, fdiff_H)

Test Passed

Пример JuMP

Код для вычисления градиента и гессиана с использованием ForwardDiff можно использовать многократно для разных операторов. Поэтому его полезно инкапсулировать в функцию:

"""
    fdiff_derivatives(f::Function) -> Tuple{Function,Function}


Return a tuple of functions that evaluate the gradient and Hessian of `f` using
ForwardDiff.jl.
"""

function fdiff_derivatives(f::Function)
    function ∇f(g::AbstractVector{T}, x::Vararg{T,N}) where {T,N}
        ForwardDiff.gradient!(g, y -> f(y...), collect(x))
        return
    end
    function ∇²f(H::AbstractMatrix{T}, x::Vararg{T,N}) where {T,N}
        h = ForwardDiff.hessian(y -> f(y...), collect(x))
        for i in 1:N, j in 1:i
            H[i, j] = h[i, j]
        end
        return
    end
    return ∇f, ∇²f
end

fdiff_derivatives (generic function with 1 method)

Вот пример использования fdiff_derivatives.

function fdiff_rosenbrock()
    model = Model(Ipopt.Optimizer)
    set_silent(model)
    @variable(model, x[1:2])
    @operator(model, op_rosenbrock, 2, f, fdiff_derivatives(f)...)
    @objective(model, Min, op_rosenbrock(x[1], x[2]))
    optimize!(model)
    Test.@test is_solved_and_feasible(model)
    return value.(x)
end

fdiff_rosenbrock()

2-element Vector{Float64}:
 0.9999999999999899
 0.9999999999999792

Enzyme

Еще одна библиотека для автоматического дифференцирования в Julia — Enzyme.jl.

Преимущества и недостатки

Главное преимущество библиотеки Enzyme в том, что она может быстро вычислять производные для функций со множеством входных аргументов.

Главные же недостатки в том, что она может выдавать необычные ошибки, если в коде используется неподдерживаемая функция Julia, и в длительной компиляции.

Разработчики JuMP не могут помочь в отладке сообщений об ошибках, связанных с Enzyme. Если оператор работает, замечательно. Если нет, мы предлагаем попробовать другую библиотеку автоматического дифференцирования. Дополнительные сведения см. на сайте juliadiff.org.

Градиент

Градиент можно вычислить с помощью Enzyme.autodiff в режиме Enzyme.Reverse. x необходимо заключить в Enzyme.Active, чтобы указать, что производные должны вычисляться по этим аргументам.

function enzyme_∇f(g::AbstractVector{T}, x::Vararg{T,N}) where {T,N}
    g .= Enzyme.autodiff(Enzyme.Reverse, f, Enzyme.Active.(x)...)[1]
    return
end

enzyme_∇f (generic function with 1 method)

Давайте проверим, получено ли аналитическое решение:

enzyme_g = zeros(2)
enzyme_∇f(enzyme_g, x...)
Test.@test ≈(analytic_g, enzyme_g)

Test Passed

Гессиан

Гессиан можно вычислить в Enzyme с использованием прямого-обратного автоматического дифференцирования.

Код для реализации гессиана в Enzyme сложен, поэтому мы не будем разбирать его подробно; см. документацию по Enzyme.

function enzyme_∇²f(H::AbstractMatrix{T}, x::Vararg{T,N}) where {T,N}
    # direction(i) возвращает кортеж с `1` в `i`-м элементе и `0`
    # в противном случае
    direction(i) = ntuple(j -> Enzyme.Active(T(i == j)), N)
    # Вычисляем градиент в обратном режиме во внутренней функции
    ∇f_deferred(x...) = Enzyme.autodiff_deferred(Enzyme.Reverse, f, x...)[1]
    # Для внешнего автоматического дифференцирования используем прямой режим.
    hess = Enzyme.autodiff(
        Enzyme.Forward,
        ∇f_deferred,
        # Проводим несколько вычислений в прямом режиме, каждый раз используя `x`, но
        # выполняя инициализацию в другом направлении.
        Enzyme.BatchDuplicated.(Enzyme.Active.(x), ntuple(direction, N))...,
    )[1]
    # Распаковываем структуру данных `hess` Enzyme в матрицу `H`, которую ожидает
    # JuMP.
    for j in 1:N, i in 1:j
        H[j, i] = hess[j][i]
    end
    return
end

enzyme_∇²f (generic function with 1 method)

Давайте проверим, получено ли аналитическое решение:

enzyme_H = zeros(2, 2)
enzyme_∇²f(enzyme_H, x...)
Test.@test ≈(analytic_H, enzyme_H)

Test Passed

Пример JuMP

Код для вычисления градиента и гессиана с использованием Enzyme можно использовать многократно для разных операторов. Поэтому его полезно инкапсулировать в функцию:

"""
    enzyme_derivatives(f::Function) -> Tuple{Function,Function}


Return a tuple of functions that evaluate the gradient and Hessian of `f` using
Enzyme.jl.
"""

function enzyme_derivatives(f::Function)
    function ∇f(g::AbstractVector{T}, x::Vararg{T,N}) where {T,N}
        g .= Enzyme.autodiff(Enzyme.Reverse, f, Enzyme.Active.(x)...)[1]
        return
    end
    function ∇²f(H::AbstractMatrix{T}, x::Vararg{T,N}) where {T,N}
        direction(i) = ntuple(j -> Enzyme.Active(T(i == j)), N)
        ∇f_deferred(x...) = Enzyme.autodiff_deferred(Enzyme.Reverse, f, x...)[1]
        hess = Enzyme.autodiff(
            Enzyme.Forward,
            ∇f_deferred,
            Enzyme.BatchDuplicated.(Enzyme.Active.(x), ntuple(direction, N))...,
        )[1]
        for j in 1:N, i in 1:j
            H[j, i] = hess[j][i]
        end
        return
    end
    return ∇f, ∇²f
end

enzyme_derivatives (generic function with 1 method)

Вот пример использования enzyme_derivatives.

function enzyme_rosenbrock()
    model = Model(Ipopt.Optimizer)
    set_silent(model)
    @variable(model, x[1:2])
    @operator(model, op_rosenbrock, 2, f, enzyme_derivatives(f)...)
    @objective(model, Min, op_rosenbrock(x[1], x[2]))
    optimize!(model)
    Test.@test is_solved_and_feasible(model)
    return value.(x)
end

enzyme_rosenbrock()

2-element Vector{Float64}:
 0.9999999999999899
 0.9999999999999792