torch.optim

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

Как использовать оптимизатор

Чтобы использовать torch.optim, необходимо создать объект оптимизатора, который будет хранить текущее состояние и обновлять параметры на основе вычисленных градиентов.

Создание оптимизатора

Для создания Optimizer необходимо передать итерируемый объект, содержащий параметры (все должны быть Variable), которые нужно оптимизировать. Затем можно указать специфичные для оптимизатора параметры, такие как скорость обучения, штраф за вес и т. д.

Пример:

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
optimizer = optim.Adam([var1, var2], lr=0.0001)

Параметры для каждого параметра

Optimizer также поддерживают указание параметров для каждого параметра. Для этого вместо передачи итерируемого объекта с Variable передайте итерируемый объект с dict s. Каждый из них определит отдельную группу параметров и должен содержать ключ params, содержащий список параметров, принадлежащих этой группе. Другие ключи должны соответствовать ключевым аргументам, принимаемым оптимизаторами, и будут использоваться в качестве параметров оптимизации для этой группы.

Например, это очень полезно, когда нужно задать скорость обучения для каждого слоя:

optim.SGD([
                {'params': model.base.parameters()},
                {'params': model.classifier.parameters(), 'lr': 1e-3}
            ], lr=1e-2, momentum=0.9)

Это означает, что параметры model.base будут использовать скорость обучения по умолчанию 1e-2, параметры model.classifier будут использовать скорость обучения 1e-3, а для всех параметров будет использоваться момент 0.9.

Выполнение шага оптимизации

Все оптимизаторы реализуют метод step(), который обновляет параметры. Его можно использовать двумя способами:

optimizer.step()

Это упрощенная версия, поддерживаемая большинством оптимизаторов. Функцию можно вызвать после вычисления градиентов, например, с помощью backward().

Пример:

for input, target in dataset:
    optimizer.zero_grad()
    output = model(input)
    loss = loss_fn(output, target)
    loss.backward()
    optimizer.step()

optimizer.step(closure)

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

Пример:

for input, target in dataset:
    def closure():
        optimizer.zero_grad()
        output = model(input)
        loss = loss_fn(output, target)
        loss.backward()
        return loss
    optimizer.step(closure)

Базовый класс

class torch.optim.Optimizer(params, defaults) [source]

Базовый класс для всех оптимизаторов.

Предупреждение

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

Parameters

• params (iterable) – итерируемый объект с torch.Tensor или dict s. Указывает, какие тензоры должны быть оптимизированы.
• defaults (Dict[str, Any]) – (dict): словарь, содержащий значения по умолчанию для параметров оптимизации (используется, когда группа параметров их не задает).

Алгоритмы

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

У нас есть 3 основные категории реализаций: цикл for, foreach (многомерный тензор) и слияние. Наиболее простые реализации представляют собой циклы for по параметрам с большими блоками вычислений. Циклы for обычно медленнее, чем наши реализации foreach, которые объединяют параметры в многомерный тензор и выполняют большие блоки вычислений сразу, тем самым экономя много последовательных вызовов ядра. Несколько наших оптимизаторов имеют ещё более быстрые объединённые реализации, которые объединяют большие блоки вычислений в одно ядро. Мы можем считать, что реализации foreach объединяют горизонтально, а объединённые — вертикально поверх этого.

В общем, порядок производительности 3 реализаций — объединённые > foreach > for-loop. Поэтому, когда это возможно, по умолчанию используется foreach вместо for-loop. Возможно, это означает, что реализация foreach доступна, пользователь не указал никаких параметров, специфичных для реализации (например, объединённые, foreach, дифференцируемые) и все тензоры являются родными и находятся на CUDA. Обратите внимание, что, хотя объединённые должны быть ещё быстрее, чем foreach, реализации являются новыми, и мы хотели бы дать им больше времени для отработки, прежде чем переключаться на них повсюду. Тем не менее, вы можете попробовать их!

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

Как настроить скорость обучения

torch.optim.lr_scheduler предоставляет несколько способов изменения скорости обучения на основе количества эпох. torch.optim.lr_scheduler.ReduceLROnPlateau позволяет динамически уменьшать скорость обучения на основе некоторых показателей валидации.

Планирование скорости обучения должно применяться после обновления оптимизатора; например, вы должны написать свой код так:

Пример:

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
scheduler = ExponentialLR(optimizer, gamma=0.9)

for epoch in range(20):
    for input, target in dataset:
        optimizer.zero_grad()
        output = model(input)
        loss = loss_fn(output, target)
        loss.backward()
        optimizer.step()
    scheduler.step()

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

Пример:

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
scheduler1 = ExponentialLR(optimizer, gamma=0.9)
scheduler2 = MultiStepLR(optimizer, milestones=[30,80], gamma=0.1)

for epoch in range(20):
    for input, target in dataset:
        optimizer.zero_grad()
        output = model(input)
        loss = loss_fn(output, target)
        loss.backward()
        optimizer.step()
    scheduler1.step()
    scheduler2.step()

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

>>> scheduler = ...
>>> for epoch in range(100):
>>>     train(...)
>>>     validate(...)
>>>     scheduler.step()

Среднее значение весов (SWA и EMA)

torch.optim.swa_utils реализует усреднение весов по методу Stochastic Weight Averaging (SWA) и экспоненциальное скользящее среднее (EMA). В частности, класс torch.optim.swa_utils.AveragedModel реализует модели SWA и EMA, torch.optim.swa_utils.SWALR реализует расписание скорости обучения SWA, а torch.optim.swa_utils.update_bn() является вспомогательной функцией, используемой для обновления статистик пакетной нормализации SWA/EMA в конце обучения.

SWA был предложен в статье Averaging Weights Leads to Wider Optima and Better Generalization.

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

Создание усреднённых моделей

Класс AveragedModel служит для вычисления весов модели SWA или EMA.

Вы можете создать усреднённую модель SWA, выполнив:

>>> averaged_model = AveragedModel(model)

Модели EMA создаются путем задания аргумента multi_avg_fn следующим образом:

>>> decay = 0.999
>>> averaged_model = AveragedModel(model, multi_avg_fn=get_ema_multi_avg_fn(decay))

Коэффициент затухания — параметр от 0 до 1, который управляет скоростью затухания усреднённых параметров. Если он не указан в get_ema_multi_avg_fn, значение по умолчанию равно 0,999.

get_ema_multi_avg_fn возвращает функцию, применяющую следующее уравнение EMA к весам:

где alpha — коэффициент затухания EMA.

Здесь модель model может быть произвольным объектом torch.nn.Module. averaged_model будет отслеживать средние значения параметров model. Для обновления этих средних значений необходимо использовать функцию update_parameters() после optimizer.step().

>>> averaged_model.update_parameters(model)

В случае SWA и EMA это вызов обычно выполняется сразу после оптимизатора step(). В случае SWA, это обычно пропускается в течение некоторого количества шагов в начале обучения.

Настраиваемые стратегии усреднения

По умолчанию torch.optim.swa_utils.AveragedModel вычисляет текущее среднее значение параметров, которые вы предоставляете, но вы также можете использовать настраиваемые функции усреднения с параметрами avg_fn или multi_avg_fn:

• avg_fn позволяет определить функцию, работающую с каждой парой параметров (усреднённый параметр, параметр модели) и должна возвращать новый усреднённый параметр.
• multi_avg_fn позволяет определить более эффективные операции, действующие на кортеже списков параметров (список усреднённых параметров, список параметров модели) одновременно, например, используя функции torch._foreach* . Эта функция должна обновлять усреднённые параметры на месте.

В следующем примере ema_model вычисляет экспоненциальное скользящее среднее с использованием параметра avg_fn:

>>> ema_avg = lambda averaged_model_parameter, model_parameter, num_averaged:\
>>>         0.9 * averaged_model_parameter + 0.1 * model_parameter
>>> ema_model = torch.optim.swa_utils.AveragedModel(model, avg_fn=ema_avg)

В следующем примере ema_model вычисляет экспоненциальное скользящее среднее с использованием более эффективного параметра multi_avg_fn:

>>> ema_model = AveragedModel(model, multi_avg_fn=get_ema_multi_avg_fn(0.9))

Расписания скорости обучения SWA

Обычно в SWA скорость обучения устанавливается на высокое постоянное значение. SWALR — это расписание скорости обучения, которое уменьшает скорость обучения до фиксированного значения, а затем поддерживает его постоянным. Например, следующий код создает расписание, которое линейно уменьшает скорость обучения от её начального значения до 0,05 за 5 эпох в каждой группе параметров:

>>> swa_scheduler = torch.optim.swa_utils.SWALR(optimizer, \
>>>         anneal_strategy="linear", anneal_epochs=5, swa_lr=0.05)

Вы также можете использовать косинусное уменьшение до фиксированного значения вместо линейного уменьшения, задав anneal_strategy="cos".

Обработка пакетной нормализации

update_bn() — это вспомогательная функция, которая позволяет вычислить статистики пакетной нормализации для модели SWA на заданном наборе данных loader в конце обучения:

>>> torch.optim.swa_utils.update_bn(loader, swa_model)

update_bn() применяет swa_model к каждому элементу в наборе данных и вычисляет статистики активации для каждого слоя пакетной нормализации в модели.

Объединение всего: SWA

В примере ниже swa_model — модель SWA, которая накапливает средние значения весов. Мы обучаем модель в течение 300 эпох и переключаемся на расписание скорости обучения SWA, а также начинаем собирать средние значения параметров SWA на 160 эпохе:

>>> loader, optimizer, model, loss_fn = ...
>>> swa_model = torch.optim.swa_utils.AveragedModel(model)
>>> scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=300)
>>> swa_start = 160
>>> swa_scheduler = SWALR(optimizer, swa_lr=0.05)
>>>
>>> for epoch in range(300):
>>>       for input, target in loader:
>>>           optimizer.zero_grad()
>>>           loss_fn(model(input), target).backward()
>>>           optimizer.step()
>>>       if epoch > swa_start:
>>>           swa_model.update_parameters(model)
>>>           swa_scheduler.step()
>>>       else:
>>>           scheduler.step()
>>>
>>> # Update bn statistics for the swa_model at the end
>>> torch.optim.swa_utils.update_bn(loader, swa_model)
>>> # Use swa_model to make predictions on test data
>>> preds = swa_model(test_input)

Объединение всего: EMA

В примере ниже ema_model — модель EMA, которая накапливает экспоненциально убывающие средние значения весов с коэффициентом затухания 0,999. Мы обучаем модель в течение 300 эпох и начинаем собирать средние значения EMA сразу.

>>> loader, optimizer, model, loss_fn = ...
>>> ema_model = torch.optim.swa_utils.AveragedModel(model, \
>>>             multi_avg_fn=torch.optim.swa_utils.get_ema_multi_avg_fn(0.999))
>>>
>>> for epoch in range(300):
>>>       for input, target in loader:
>>>           optimizer.zero_grad()
>>>           loss_fn(model(input), target).backward()
>>>           optimizer.step()
>>>           ema_model.update_parameters(model)
>>>
>>> # Update bn statistics for the ema_model at the end
>>> torch.optim.swa_utils.update_bn(loader, ema_model)
>>> # Use ema_model to make predictions on test data
>>> preds = ema_model(test_input)