Документация Engee

Принципы работы Flux: градиенты и слои

Взятие градиентов

Основные функции Flux используют код Julia для взятия градиентов. Функция gradient принимает другую функцию Julia f и набор аргументов и возвращает градиент относительно каждого аргумента. (Попробуйте вставить код этих примеров в терминале Julia.)

julia> using Flux

julia> f(x) = 3x^2 + 2x + 1;

julia> df(x) = gradient(f, x)[1]; # df/dx = 6x + 2

julia> df(2)
14.0

julia> d2f(x) = gradient(df, x)[1]; # d²f/dx² = 6

julia> d2f(2)
6.0

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

julia> f(x, y) = sum((x .- y).^2);

julia> gradient(f, [2, 1], [2, 0])
([0.0, 2.0], [-0.0, -2.0])

Эти градиенты основаны на x и y. Вместо этого во Flux градиенты берутся на основе весов и смещений, составляющих параметры модели.

В машинном обучении нередко приходится сталкиваться с сотнями массивов параметров. Вместо передачи их в gradient по отдельности их можно хранить вместе в одной структуре. Простейшим примером является именованный кортеж, создаваемый с помощью такого синтаксиса:

julia> nt = (a = [2, 1], b = [2, 0], c = tanh);

julia> g(x::NamedTuple) = sum(abs2, x.a .- x.b);

julia> g(nt)
1

julia> dg_nt = gradient(g, nt)[1]
(a = [0.0, 2.0], b = [-0.0, -2.0], c = nothing)

Обратите внимание, что функция gradient вернула соответствующую структуру. Поле dg_nt.a является градиентом для nt.a и т. д. У некоторых полей нет градиента, на что указывает значение nothing.

Вместо того чтобы каждый раз определять функцию наподобие g (и придумывать ей имя), часто бывает удобнее использовать анонимные функции, в данном случае x -> sum(abs2, x.a .- x.b). Анонимную функцию можно определить с помощью оператора -> или do, причем блоки do полезны, если нужно выполнить несколько действий:

julia> gradient((x, y) -> sum(abs2, x.a ./ y .- x.b), nt, [1, 2])
((a = [0.0, 0.5], b = [-0.0, -1.0], c = nothing), [-0.0, -0.25])

julia> gradient(nt, [1, 2]) do x, y
         z = x.a ./ y
         sum(abs2, z .- x.b)
       end
((a = [0.0, 0.5], b = [-0.0, -1.0], c = nothing), [-0.0, -0.25])

Иногда необходимо знать значение функции, а также ее градиент. Вместо того чтобы вызывать функцию еще раз, можно вызвать withgradient:

julia> Flux.withgradient(g, nt)
(val = 1, grad = ((a = [0.0, 2.0], b = [-0.0, -2.0], c = nothing),))

!!! note "Implicit gradients" Ранее во Flux многие параметры обрабатывались иначе с помощью функции params. При этом используется метод gradient, принимающий функцию без аргументов и возвращающий словарь, через который можно искать итоговые градиенты:

```julia-repl
julia> x = [2, 1];

julia> y = [2, 0];

julia> gs = gradient(Flux.params(x, y)) do
         f(x, y)
       end
Grads(...)

julia> gs[x]
2-element Vector{Float64}:
 0.0
 2.0

julia> gs[y]
2-element Vector{Float64}:
 -0.0
 -2.0
```

Создание простых моделей

Рассмотрим простую линейную регрессию, которая пытается спрогнозировать выходной массив y на основе входных данных x.

W = rand(2, 5)
b = rand(2)

predict(x) = W*x .+ b

function loss(x, y)
  ŷ = predict(x)
  sum((y .- ŷ).^2)
end

x, y = rand(5), rand(2) # Фиктивные данные
loss(x, y) # ~ 3

Чтобы улучшить результат прогнозирования, мы можем взять градиенты потерь относительно W и b и выполнить градиентный спуск.

using Flux

gs = gradient(() -> loss(x, y), Flux.params(W, b))

Теперь, когда у нас есть градиенты, мы можем их извлечь и обновить W для обучения модели.

W̄ = gs[W]

W .-= 0.1 .* W̄

loss(x, y) # ~ 2,5

Потери немного уменьшились, а значит, наш прогноз x стал ближе к целевому значению y. При наличии некоторых данных уже можно попробовать обучить модель.

Любое глубокое обучение во Flux независимо от уровня сложности является частным случаем этого примера. Безусловно, модели могут выглядеть совершенно иначе — в них могут быть миллионы параметров или сложный порядок выполнения. Давайте посмотрим, как Flux работает с более сложными моделями.

Построение слоев

Обычно модели бывают гораздо сложнее, чем в приведенном выше примере линейной регрессии. Например, это могут быть два линейных слоя с нелинейной функцией между ними, например сигмоидом (σ). Если следовать приведенному выше стилю, это можно записать так:

using Flux

W1 = rand(3, 5)
b1 = rand(3)
layer1(x) = W1 * x .+ b1

W2 = rand(2, 3)
b2 = rand(2)
layer2(x) = W2 * x .+ b2

model(x) = layer2(σ.(layer1(x)))

model(rand(5)) # => 2-элементный вектор

Такой вариант сработает, но он слишком громоздкий — в нем много повторов, а будет еще больше, если добавить дополнительные слои. Один из способов факторизации заключается в создании функции, которая возвращает линейные слои.

function linear(in, out)
  W = randn(out, in)
  b = randn(out)
  x -> W * x .+ b
end

linear1 = linear(5, 3) # можно обратиться к linear1.W и т. д.
linear2 = linear(3, 2)

model(x) = linear2(σ.(linear1(x)))

model(rand(5)) # => 2-элементный вектор

Другим (равносильным) способом является создание структуры, явным образом представляющей аффинный слой.

struct Affine
  W
  b
end

Affine(in::Integer, out::Integer) =
  Affine(randn(out, in), randn(out))

# Перегружаем вызов, чтобы объект можно было использовать как функцию
(m::Affine)(x) = m.W * x .+ m.b

a = Affine(10, 5)

a(rand(10)) # => 5-элементный вектор

Поздравляем! Вы только что построили слой Dense, который есть во Flux. Во Flux доступно множество интересных слоев, но вы могли бы очень легко реализовать их сами.

(В случае с Dense есть одно небольшое отличие: для удобства этот слой принимает также функцию активации, например Dense(10 => 5, σ).)

Составление очередности

Достаточно часто приходится создавать модели наподобие следующей:

layer1 = Dense(10 => 5, σ)
# ...
model(x) = layer3(layer2(layer1(x)))

Для длинных цепочек было бы естественнее составить список слоев, например так:

using Flux

layers = [Dense(10 => 5, σ), Dense(5 => 2), softmax]

model(x) = foldl((x, m) -> m(x), layers, init = x)

model(rand(10)) # => 2-элементный вектор

К счастью, во Flux есть и такая возможность:

model2 = Chain(
  Dense(10 => 5, σ),
  Dense(5 => 2),
  softmax)

model2(rand(10)) # => 2-элементный вектор

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

Приятной особенностью такого подхода является то, что поскольку «модели» — это просто функции (возможно, с обучаемыми параметрами), данный механизм можно представить как обычную композицию функций.

m = Dense(5 => 2) ∘ Dense(10 => 5, σ)

m(rand(10))

Аналогичным образом, Chain будет спокойно работать с любой функцией Julia.

m = Chain(x -> x^2, x -> x+1)

m(5) # => 26

Вспомогательные функции слоев

С данным слоем Affine остается одна проблема: Flux не может заглянуть внутрь него. Это означает, что функции Flux.train! будут недоступны его параметры, а функция gpu не сможет переместить его на GPU. Такие возможности обеспечиваются макросом @functor:

Flux.@functor Affine

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

function Affine((in, out)::Pair; bias=true, init=Flux.randn32)
  W = init(out, in)
  b = Flux.create_bias(W, bias, out)
  Affine(W, b)
end

Affine(3 => 1, bias=false, init=ones) |> gpu