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 с параметрами. Каждый из них определит отдельную группу параметров и должен содержать ключ 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(), который обновляет параметры. Он может использоваться двумя способами:

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

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)

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

Класс AveragedModel служит для вычисления весов модели SWA. Вы можете создать усредненную модель, выполнив следующие действия:

>>> swa_model = AveragedModel(model)

Здесь модель model может быть любым объектом torch.nn.Module. swa_model будет отслеживать текущие средние значения параметров model. Чтобы обновить эти средние значения, вы можете использовать функцию update_parameters():

>>> swa_model.update_parameters(model)

Планировщики скорости обучения 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 к каждому элементу в наборе данных и вычисляет статистику активации для каждого слоя пакетной нормализации в модели.

Настройка стратегий усреднения

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

Пример:

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

Объединение всего в одно целое

В примере ниже, 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)