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

Определение настраиваемых слоев

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

Пример пользовательской модели

Вот простейший пример пользовательской модели. В нем просто добавляется вход в результат работы нейронной сети.

struct CustomModel
  chain::Chain
end

function (m::CustomModel)(x)
  # Здесь может быть произвольный код, но имейте в виду, что все будет дифференцироваться.
  # В Zygote некоторые операции, например изменение массивов, невозможны.

  return m.chain(x) + x
end

# Для обучения вызовите @functor. Подробнее см. ниже.
Flux.@functor CustomModel

После этого модель можно использовать следующим образом.

chain = Chain(Dense(10, 10))
model = CustomModel(chain)
model(rand(10))

Введение во Flux и автоматическое дифференцирование см. в этом руководстве.

Настройка сбора параметров для модели

Здесь мы будем ссылаться на пример слоя Affine из раздела с основами.

По умолчанию в качестве параметров собираются все поля типа Affine, однако иногда в «слоях» могут содержаться другие метаданные, которые не требуются для обучения и поэтому должны игнорироваться при сборе параметров. Во Flux поля слоя помечаются как обучаемые путем перегрузки функции trainable:

julia> Flux.@functor Affine

julia> a = Affine(Float32[1 2; 3 4; 5 6], Float32[7, 8, 9])
Affine(Float32[1.0 2.0; 3.0 4.0; 5.0 6.0], Float32[7.0, 8.0, 9.0])

julia> Flux.params(a) # поведение по умолчанию
Params([Float32[1.0 2.0; 3.0 4.0; 5.0 6.0], Float32[7.0, 8.0, 9.0]])

julia> Flux.trainable(a::Affine) = (; a.W)  # возвращает NamedTuple с использованием имени поля

julia> Flux.params(a)
Params([Float32[1.0 2.0; 3.0 4.0; 5.0 6.0]])

При вызове Flux.params только поля, возвращаемые функцией trainable, будут собираться как обучаемые параметры, и только они будут доступны функциям Flux.setup и Flux.update! для обучения. Однако функции gpu и другим аналогичным функциям будут доступны все поля, например:

julia> a |> f16
Affine(Float16[1.0 2.0; 3.0 4.0; 5.0 6.0], Float16[7.0, 8.0, 9.0])

Имейте в виду, что для скрытия полей, которые не содержат обучаемых параметров (например, функций активации или логических флагов), перегружать trainable не нужно. Они всегда игнорируются params и при обучении:

julia> Flux.params(Affine(true, [10, 11, 12.0]))
Params([])

Доступные поля можно ограничить дополнительно с помощью выражения @functor Affine (W,). Однако делать это не рекомендуется. Для этого требуется, чтобы структура struct имела соответствующий конструктор, принимающий только W в качестве аргумента, и игнорируемые поля будут недоступны таким функциям, как gpu (что обычно нежелательно).

Фиксация параметров слоя

Если нужно включить не все параметры модели (например, для переноса обучения), можно просто не передавать эти слои в вызов params.

Совместимость: Flux ≤ 0.14

Описываемый здесь механизм предназначен для прежнего «неявного» стиля обучения во Flux. При обновлении до Flux 0.15 его следует заменить freeze! и thaw!.

Рассмотрим простую многослойную модель перцептрона, в которой не следует оптимизировать первые два слоя Dense. Этого можно добиться с помощью функций срезов, предоставляемых Chain:

m = Chain(
      Dense(784 => 64, relu),
      Dense(64 => 64, relu),
      Dense(32 => 10)
    );

ps = Flux.params(m[3:end])

Объект ps типа Zygote.Params теперь содержит ссылку только на переданные в него параметры слоев.

Во время обучения градиенты будут вычисляться только для последнего слоя Dense (и применяться только к нему), поэтому изменятся только его параметры.

Flux.params также принимает несколько входов для упрощения сбора параметров из разнородных моделей одним вызовом. Это можно легко продемонстрировать на основе предыдущего примера, в котором нужно исключить оптимизацию второго слоя Dense. Выглядеть это будет примерно так.

Flux.params(m[1], m[3:end])

Иногда требуется более детальный контроль. Мы можем зафиксировать определенный параметр слоя, который уже введен в объект ps типа Params, просто удалив его из ps:

ps = Flux.params(m)
delete!(ps, m[2].bias)

Пользовательский слой с несколькими входами или выходами

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

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

Несколько входов: пользовательский слой Join

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

Сначала определим новую структуру Join, в полях которой будут храниться различные пути и операция объединения.

using Flux
using CUDA

# пользовательский слой объединения
struct Join{T, F}
  combine::F
  paths::T
end

# конструктором может быть Join(op, m1, m2, ...)
Join(combine, paths...) = Join(combine, paths)

Обратите внимание, что мы параметризовали тип поля paths. Это необходимо для ускорения выполнения кода Julia; в общем случае T может быть Tuple или Vector, но обращать на это внимание не требуется. То же самое касается поля combine.

Далее необходимо использовать Functors.@functor, чтобы структура работала как слой Flux. Это необходимо для того, чтобы при вызове params для Join возвращались базовые массивы весов для каждого пути.

Flux.@functor Join

Наконец, определим прямой проход. Для Join это означает применение каждого path в paths к каждому входному массиву, а затем использование combine для слияния результатов.

(m::Join)(xs::Tuple) = m.combine(map((f, x) -> f(x), m.paths, xs)...)
(m::Join)(xs...) = m(xs)

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

model = Chain(
              Join(vcat,
                   Chain(Dense(1 => 5, relu), Dense(5 => 1)), # ветвь 1
                   Dense(1 => 2),                             # ветвь 2
                   Dense(1 => 1)                              # ветвь 3
                  ),
              Dense(4 => 1)
             ) |> gpu

xs = map(gpu, (rand(1), rand(1), rand(1)))

model(xs)
# возвращает один вектор с плавающей запятой, содержащий одно значение

Слой Join доступен в пакете Fluxperimental.jl.

Использование Parallel

Во Flux уже есть тип Parallel, предоставляющую такую функциональность. В данном случае Join будет просто "синтаксическим сахаром" для Parallel.

Join(combine, paths) = Parallel(combine, paths)
Join(combine, paths...) = Join(combine, paths)

# используем версию с переменным числом аргументов или кортежем для прямого прохода Parallel
model = Chain(
              Join(vcat,
                   Chain(Dense(1 => 5, relu), Dense(5 => 1)),
                   Dense(1 => 2),
                   Dense(1 => 1)
                  ),
              Dense(4 => 1)
             ) |> gpu

xs = map(gpu, (rand(1), rand(1), rand(1)))

model(xs)
# возвращает один вектор с плавающей запятой, содержащий одно значение

Несколько выходов: пользовательский слой Split

Наш пользовательский слой Split будет принимать один вход, а затем передавать его по отдельному пути для получения нескольких выходов.

Начнем с тех же действий, что и в случае со слоем Join: определим структуру, используем Functors.@functor, а затем определим прямой проход.

using Flux
using CUDA

# пользовательский слой разделения
struct Split{T}
  paths::T
end

Split(paths...) = Split(paths)

Flux.@functor Split

(m::Split)(x::AbstractArray) = map(f -> f(x), m.paths)

Теперь можно проверить, действительно ли Split выдает несколько выходов.

model = Chain(
              Dense(10 => 5),
              Split(Dense(5 => 1, tanh), Dense(5 => 3, tanh), Dense(5 => 2))
             ) |> gpu

model(gpu(rand(10)))
# возвращает кортеж с тремя векторами чисел с плавающей запятой

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

using Statistics

# предполагаем, что модель возвращает выход Split
# x — единственный вход
# ys — кортеж выходов
function loss(x, ys, model)
  # rms по всем mse
  ŷs = model(x)
  return sqrt(mean(Flux.mse(y, ŷ) for (y, ŷ) in zip(ys, ŷs)))
end

Слой Split доступен в пакете Fluxperimental.jl.