Советы по производительности

В этом разделе кратко описаны некоторые общие рекомендации по достижению хорошей производительности при использовании Turing. Общие методы по обеспечению надлежащей производительности программ Julia приводятся в документации по Julia.

Используйте многомерные распределения

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

Следующий пример:

@model function gmodel(x)
    m ~ Normal()
    for i = 1:length(x)
        x[i] ~ Normal(m, 0.2)
    end
end

можно выразить более эффективно с помощью простого преобразования:

using FillArrays

@model function gmodel(x)
    m ~ Normal()
    x ~ MvNormal(Fill(m, length(x)), 0.04 * I)
end

Выбирайте соответствующий бэкенд AD

В настоящее время Turing поддерживает два различных бэкенда автоматического дифференцирования (AD). Как правило, старайтесь использовать :forwarddiff для моделей с небольшим количеством параметров и :reversediff, :tracker или :zygote для моделей с большими векторами параметров или операциями линейной алгебры. Подробные сведения см. в статье об автоматическом дифференцировании.

Уделяйте особое внимание `:tracker` и `:zygote`

В случае с :tracker и :zygote необходимо пока избегать циклов. В основном это связано с обратным режимом работы AD в бэкендах Tracker и Zygote, которые неэффективны для таких ситуаций. ReverseDiff работает лучше, но векторизированные операции все равно будут выполняться эффективнее.

Избежать циклов можно с помощью filldist(dist, N) и arraydist(dists). filldist(dist, N) создает многомерное распределение, состоящее из N одинаковых и независимых копий одномерного распределения dist, если dist является одномерным, или создает матрично-вариативное распределение, состоящее из N одинаковых и независимых копий многомерного распределения dist, если dist является многомерным. filldist(dist, N, M) также можно использовать для создания матрично-вариативного распределения из одномерного распределения dist. arraydist(dists) аналогичен filldist, но в качестве входных данных принимает массив распределений dists. Написание пользовательского распределения с пользовательским сопряжением является еще одним вариантом отказа от циклов.

Обеспечьте возможность вывода типов в модели

Для эффективного вывода на основе градиента, например при использовании HMC, NUTS или ADVI, важно обеспечить возможность вывода типов в вашей модели.

Следующий пример с абстрактными типами

@model function tmodel(x, y)
    p,n = size(x)
    params = Vector{Real}(undef, n)
    for i = 1:n
        params[i] ~ truncated(Normal(), 0, Inf)
    end

    a = x * params
    y ~ MvNormal(a, I)
end

может быть преобразовано в следующее представление с конкретными типами:

@model function tmodel(x, y, ::Type{T} = Float64) where {T}
    p,n = size(x)
    params = Vector{T}(undef, n)
    for i = 1:n
        params[i] ~ truncated(Normal(), 0, Inf)
    end

    a = x * params
    y ~ MvNormal(a, I)
end

В качестве альтернативы в этом примере можно использовать filldist:

@model function tmodel(x, y)
    params ~ filldist(truncated(Normal(), 0, Inf), size(x, 2))
    a = x * params
    y ~ MvNormal(a, I)
end

Обратите внимание, что вы можете использовать @code_warntype для поиска типов в определении модели, которые компилятор не может вывести. В REPL Julia они помечены красным цветом.

Например, рассмотрим следующую простую программу:

@model function tmodel(x)
    p = Vector{Real}(undef, 1)
    p[1] ~ Normal()
    p = p .+ 1
    x ~ Normal(p[1])
end

Мы можем использовать

using Random

model = tmodel(1.0)

@code_warntype model.f(
    model,
    Turing.VarInfo(model),
    Turing.SamplingContext(
        Random.default_rng(), Turing.SampleFromPrior(), Turing.DefaultContext(),
    ),
    model.args...,
)

для проверки вывода типов в модели.

Повторно используйте вычисления в выборке Гиббса

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

@model function demo(x)
    a ~ Gamma()
    b ~ Normal()
    c = function1(a)
    d = function2(b)
    x .~ Normal(c, d)
end
alg = Gibbs(MH(:a), MH(:b))
sample(demo(zeros(10)), alg, 1000)

если в субитерации Гиббса обновлять только a, сохраняя b неизменным, то значение d не меняется. И если обновлять только b, значение c не меняется. Однако если function1 и function2 являются дорогостоящими и выполняются в каждой субитерации Гиббса, на вычисление значений, которые мы уже вычислили ранее, уйдет много времени. Такую проблему можно решить с помощью Memoization.jl. Мемоизация функции позволяет хранить и повторно использовать вывод функции для каждого входного значения, с которым она вызывается. Это связано с небольшими временными затратами, но для дорогих функций экономия будет намного больше.

Чтобы использовать Memoization.jl, просто определите мемоизированные версии function1 и function2 следующим образом:

using Memoization

@memoize memoized_function1(args...) = function1(args...)
@memoize memoized_function2(args...) = function2(args...)

Затем определите модель Turing, используя новые функции, следующим образом:

@model function demo(x)
    a ~ Gamma()
    b ~ Normal()
    c = memoized_function1(a)
    d = memoized_function2(b)
    x .~ Normal(c, d)
end