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

Применяются все обычные советы по повышению производительности Julia. Как всегда, профилирование кода является полезным способом поиска узких мест. Ниже приведены некоторые характерные для Flux советы и напоминания.

Не используйте большую точность, чем нужно

Flux отлично работает со всеми числовыми типами. Но часто не требуется работать, скажем, с Float64 (не говоря уже о BigFloat). Переход на Float32 может значительно повысить скорость, но не потому, что операции выполняются быстрее, а потому, что использование памяти сокращается вдвое. Это означает, что выделение происходит гораздо быстрее. И вы используете меньший объем памяти.

Сохраняйте типы входных данных

Функции активации и потерь не только должны быть устойчивыми по типу, они также должны сохранять тип своих входных данных.

Искусственный пример с использованием функции активации типа

my_tanh(x) = Float64(tanh(x))

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

Это означает, что если вы измените данные, скажем, с Float64 на Float32 (что должно привести к ускорению: см. выше), вы увидите значительное замедление.

Это может произойти незаметно, потому что вы можете вызвать продвижение типа, взаимодействуя с числовыми литералами. Например, в следующем случае вы столкнетесь с той же проблемой, что и в предыдущем:

leaky_tanh(x) = 0.01*x + tanh(x)

Хотя можно изменить функцию активации (например, использовать 0.01f0*x), идиоматическим (и безопасным) способом избежать приведения типов при изменении входных данных является использование oftype:

leaky_tanh(x) = oftype(x/1, 0.01)*x + tanh(x)

Определяйте пакеты как матрицы признаков

Иногда может возникнуть соблазн обрабатывать наблюдения (векторы признаков) по одному, например

function loss_total(xs::AbstractVector{<:Vector}, ys::AbstractVector{<:Vector})
    sum(zip(xs, ys)) do (x, y_target)
        y_pred = model(x)  # определение модели
        return loss(y_pred, y_target)
    end
end

Однако гораздо быстрее объединить их в матрицу, так как это приведет к умножению матрицы на матрицу BLAS, что намного быстрее, чем эквивалентная последовательность матрично-векторных умножений. Улучшение достаточно велико, поэтому стоит выделить новую память для хранения в непрерывной последовательности.

x_batch = reduce(hcat, xs)
y_batch = reduce(hcat, ys)
...
function loss_total(x_batch::Matrix, y_batch::Matrix)
    y_preds = model(x_batch)
    sum(loss.(y_preds, y_batch))
end

При такой конкатенации используйте reduce(hcat, xs), а не hcat(xs...). Это позволит избежать ухудшения при разделении, и вы сможете использовать оптимизированный метод reduce.