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

Обучение модели Flux

Под обучением понимается процесс медленной настройки параметров модели для улучшения ее работы. Кроме самой модели, нам понадобятся три вещи.

  • Целевая функция, которая оценивает, насколько хорошо работает модель на основе некоторых входных данных.

  • Правило оптимизации, которое описывает, как следует настраивать параметры модели.

  • Некоторые обучающие данные для использования в качестве входных данных в этом процессе.

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

# Инициализируем оптимизатор для этой модели:
opt_state = Flux.setup(rule, model)

for data in train_set
  # Распакуем этот элемент (для контролируемого обучения):
  input, label = data

  # Вычислим градиент цели
  # с учетом параметров в модели:
  grads = Flux.gradient(model) do m
      result = m(input)
      loss(result, label)
  end

  # Обновим параметры так, чтобы уменьшить цель,
  # в соответствии с выбранным правилом оптимизации:
  Flux.update!(opt_state, model, grads[1])
end

Этот цикл также можно записать с помощью функции train!, но сначала полезно разобраться в его деталях:

train!(model, train_set, opt_state) do m, x, y
  loss(m(x), y)
end

Градиенты модели

Сначала вспомните из раздела о работе с градиентами, что Flux.gradient(f, a, b) всегда вызывает f(a, b), а возвращает кортеж (∂f_∂a, ∂f_∂b). В приведенном выше коде функция f, переданная градиенту (gradient), является анонимной функцией с одним аргументом, созданной блоком do, поэтому grads — это кортеж с одним элементом. Вместо блока do можно было бы написать следующее.

grads = Flux.gradient(m -> loss(m(input), label), model)

Поскольку модель представляет собой некий вложенный набор слоев, grads[1] является аналогичным вложенным набором NamedTuple, в конечном итоге содержащим компоненты градиента. Если, например, θ = model.layers[1].weight[2,3] — это один скалярный параметр, запись в матрице весов, то производная потерь по нему равна ∂f_∂θ = grads[1].layers[1].weight[2,3].

Важно, чтобы выполнение модели происходило внутри вызова gradient, чтобы влияние параметров модели было замечено Zygote.

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

Совместимость: Implicit gradients

Во Flux версии не выше 0.14 использовался «неявный» режим Zygote, в котором gradient принимает функцию с нулевым аргументом. Он выглядит следующим образом:

pars = Flux.params(model)
grad = gradient(() -> loss(model(input), label), pars)

Здесь pars::Params и grad::Grads являются двумя подобными словарю структурам. Поддержка этой возможности будет удалена в версии Flux 0.15. В этих сине-голубых блоках объясняется, что нужно изменить.

Функции потерь

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

Это число может быть получено с помощью любого обычного кода Julia, но он должен быть выполнен внутри вызова gradient. Например, можно определить функцию

loss(y_hat, y) = sum((y_hat .- y).^2)

или написать это прямо внутри блока do выше. Многие часто используемые функции, такие как mse для среднеквадратичной погрешности или crossentropy для потерь перекрестной энтропии, доступны из модуля Flux.Losses.

Совместимость: Implicit-style loss functions

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

loss(x, y) = sum((model(x) .- y).^2)

что определяет функцию, ссылающуюся на определенную глобальную переменную model.

Правила оптимизации

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

Для градиентного спуска требуется скорость обучения, которая представляет собой небольшое число, описывающее скорость спуска, обычно записываемое в виде греческой буквы «эта» (η). Его часто называют гиперпараметром, чтобы отличать от обновляемых параметров θ = θ - η * ∂loss_∂θ. Мы хотим обновить все параметры в модели следующим образом:

η = 0.01   # скорость обучения

# Для каждого массива параметров обновим
# в соответствии с нужным градиентом:
fmap(model, grads[1]) do p, g
  p .= p .- η .* g
end

Более уточненная версия этого цикла для обновления всех параметров заключена в оболочку в виде функции (opt_state, model, grads" class="xref page">1). И скорость обучения — это единственное, что хранится в структуре Descent.

Однако существует множество других правил оптимизации, которые регулируют размер и направление шага различными продуманными способами. Для большинства из них требуется запоминать градиенты с предыдущих шагов, а не всегда выполнять прямой спуск — Momentum является самым простым. Функция setup создает необходимое для этого хранилище для конкретной модели. Она вызывается один раз перед началом обучения и возвращает древовидный объект, который является первым аргументом update!. Это выглядит следующим образом:

# Инициализируем импульс
opt_state = Flux.setup(Momentum(0.01, 0.9), model)

for data in train_set
  grads = [...]

  # Обновим параметры модели и состояние оптимизатора:
  Flux.update!(opt_state, model, grads[1])
end

Многие часто используемые правила оптимизации, такие как Adam, являются встроенными. Они перечислены на странице, посвященной оптимизаторам.

Совместимость: Implicit-style optimiser state

Эта функция setup создает еще одну древовидную структуру. В старых версиях Flux этого не происходило. Там подобная словарю структура хранилась внутри оптимизатора Adam(0.001). Он инициализировался при первом использовании версии update! для «неявных» параметров.

Наборы данных и пакеты

Приведенный выше цикл выполняет итерацию по train_set, ожидая на каждом шаге кортеж (input, label). Таким простейшим объектом является вектор кортежей, например следующий:

x = randn(28, 28)
y = rand(10)
data = [(x, y)]

или data = [(x, y), (x, y), (x, y)] для одних и тех же значений три раза.

Очень часто начальные данные представляют собой большие массивы, которые необходимо разделить на примеры. Чтобы создать один итератор пар (x, y), вам может понадобиться zip:

X = rand(28, 28, 60_000);  # много изображений, каждое размером 28 × 28
Y = rand(10, 60_000)
data = zip(eachslice(X; dims=3), eachcol(Y))

first(data) isa Tuple{AbstractMatrix, AbstractVector}  # верно

Здесь каждая итерация будет использовать одну матрицу x (возможно, изображение) и один вектор y. Очень часто вместо этого обучение проводится на основе пакетов таких входных данных (или мини-пакетов, эти два слова означают одно и то же) как для эффективности, так и для получения лучших результатов. Это легко сделать с помощью DataLoader:

data = Flux.DataLoader((X, Y), batchsize=32)

x1, y1 = first(data)
size(x1) == (28, 28, 32)
length(data) == 1875 === 60_000 ÷ 32

Слои Flux настроены для приема такого пакета входных данных, и он требуется для сверточных слоев, таких как Conv. Индекс пакета всегда является последним измерением.

Обучающие циклы

Простые обучающие циклы, подобные приведенному выше, можно компактно записать с помощью функции train!. Включая setup, это читается так:

opt_state = Flux.setup(Adam(), model)

for epoch in 1:100
  Flux.train!(model, train_set, opt_state) do m, x, y
    loss(m(x), y)
  end
end

Или можно явно написать анонимную функцию, которую создает этот блок do. train!((m,x,y) -> loss(m(x),y), model, train_set, opt_state) является точным эквивалентом.

Совместимость: Implicit-style train!

Это новый метод функции train!, который принимает результат функции setup в качестве 4-го аргумента. Первый аргумент — это функция, которая принимает саму модель. В версиях Flux не выше 0.14 существует метод train! для «неявных» параметров, который работает следующим образом:

 
```
train!((x,y) -> loss(model(x), y), Flux.params(model), train_set, Adam())
```

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

  • Функция withgradient аналогична gradient, но также возвращает значение функции для регистрации в журнале или диагностики.

  • Запись в журнал или вывод лучше всего выполнять вне вызова gradient, так как нет необходимости различать эти команды.

  • Для выхода из частей цикла используются ключевые слова Julia break и continue.

opt_state = Flux.setup(Adam(), model)

my_log = []
for epoch in 1:100
  losses = Float32[]
  for (i, data) in enumerate(train_set)
    input, label = data

    val, grads = Flux.withgradient(model) do m
      # Любой код здесь дифференцирован.
      # Оценка модели и потери должны быть выполнены внутри.
      result = m(input)
      my_loss(result, label)
    end

    # Сохраняем потери от прямого прохода. (Выполняется вне градиента.)
    push!(losses, val)

    # Обнаруживаем потери Inf или NaN. Выводим предупреждение, а затем пропускаем обновления!
    if !isfinite(val)
      @warn "loss is $val on item $i" epoch
      continue
    end

    Flux.update!(opt_state, model, grads[1])
  end

  # Вычисляем точность и сохраняем данные в виде NamedTuple
  acc = my_accuracy(model, train_set)
  push!(my_log, (; acc, losses))

  # Останавливаем обучение при достижении некоторого критерия
  if  acc > 0.95
    println("stopping after $epoch epochs")
    break
  end
end

Регуляризация

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

Некоторые из них можно реализовать, просто изменив функцию потерь. Регуляризация L₂ (иногда называемая гребневой регрессией) добавляет к потерям штраф, пропорциональный θ^2 для каждого скалярного параметра. Для очень простой модели это может быть реализовано следующим образом.

grads = Flux.gradient(densemodel) do m
  result = m(input)
  penalty = sum(abs2, m.weight)/2 + sum(abs2, m.bias)/2
  my_loss(result, label) + 0.42 * penalty
end

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

pen_l2(x::AbstractArray) = sum(abs2, x)/2

grads = Flux.gradient(model) do m
  result = m(input)
  penalty = sum(pen_l2, Flux.params(m))
  my_loss(result, label) + 0.42 * penalty
end

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

opt_state = Flux.setup(Adam(0.1), model)

этим:

decay_opt_state = Flux.setup(OptimiserChain(WeightDecay(0.42), Adam(0.1)), model)

Оптимизаторы Flux на самом деле являются изменениями, применяемыми к градиенту перед его использованием для обновления параметров, и OptimiserChain применяет два таких изменения. Первое, WeightDecay, добавляет к градиенту увеличенный в 0.42 раза исходный параметр, совпадающий с градиентом штрафа выше (с той же нереально большой константой). После этого в любом случае Adam вычисляет окончательное обновление.

Тот же механизм OptimiserChain можно использовать и для других целей, например для отсечения градиентов с помощью ClipGrad или ClipNorm.

Помимо L2 / снижения весов, другой распространенный и совершенно иной вид регуляризации доступен в слое Dropout. Он отключает некоторые выводы предыдущего слоя во время обучения. Он должен переключаться автоматически, но см. описание trainmode! / testmode! со сведениями о включении или отключении этого слоя вручную.

Замораживание и расписания

Для более детализированного управления обучением может потребоваться изменить скорость обучения в середине процесса. Это можно сделать с помощью adjust! следующим образом:

opt_state = Flux.setup(Adam(0.1), model)  # инициализируем один раз

for epoch in 1:1000
  train!([...], state)  # Обучаем при η = 0,1 для первых 100,
  if epoch == 100       # затем изменяем на использование η = 0,01 для оставшейся части.
    Flux.adjust!(opt_state, 0.01)
  end
end
Совместимость: Flux ≤ 0.14

В старом «неявном» оптимизаторе opt = Adam(0.1) эквивалентом было прямое изменение структуры Adam opt.eta = 0.001.

Можно настроить и другие гиперпараметры, например Flux.adjust!(opt_state, beta = (0.8, 0.99)). Причем такие изменения быть применены только к одной части модели. Например, задается разная скорость обучения для кодера и декодера:

# Рассмотрим модель с двумя частями:
bimodel = Chain(enc = [...], dec = [...])

# Возвращается дерево, структура которого соответствует модели:
opt_state = Flux.setup(Adam(0.02), bimodel)

# Настраиваем скорость обучения, которая будет использоваться для bimodel.layers.enc
Flux.adjust!(opt_state.layers.enc, 0.03)

Чтобы полностью отключить обучение некоторой части модели, используйте freeze!. Это временное изменение, отменяемое функцией thaw!:

Flux.freeze!(opt_state.layers.enc)

# Теперь обучение не будет обновлять параметры в bimodel.layers.enc
train!(loss, bimodel, data, opt_state)

# Разморозим всю модель:
Flux.thaw!(opt_state)
Совместимость: Flux ≤ 0.14

Ранее «неявный» эквивалент заключался в передаче в gradient объекта, ссылающегося только на часть модели, например Flux.params(bimodel.layers.enc).

Явный или неявный?

Раньше работа с градиентами, обучением и правилами оптимизации во Flux происходила совсем по-другому. Новый стиль, описанный выше, в Zygote называется «явным», а старый — «неявным». Flux 0.13 и 0.14 являются переходными версиями, которые поддерживают оба варианта.

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

Подробнее о двух режимах градиента см. в документации по Zygote.