Класс CosineAnnealingLR

Класс CosineAnnealingLR представляет собой планировщик скорости обучения, который изменяет скорость обучения по закону косинуса. Он используется вместе с оптимизатором и позволяет плавно уменьшать скорость обучения от начального значения до минимального в течение заданного количества эпох. Первый параметр класса optimizer принимает экземпляр оптимизатора, второй параметр T_max задаёт количество эпох, за которое скорость обучения должна достигнуть минимума. Также класс принимает необязательные параметры: eta_min задаёт минимальное значение скорости обучения, а last_epoch позволяет продолжить обучение с определённой эпохи.

Синтаксис

torch.optim.lr_scheduler.CosineAnnealingLR(
	optimizer,
	T_max,
	eta_min=0,
	last_epoch=-1
)

Пример

Рассмотрим базовый пример использования класса CosineAnnealingLR. Создадим простую модель и оптимизатор, а затем будем обновлять скорость обучения каждую эпоху:

import torch
import torch.nn as nn
from torch.optim import SGD
from torch.optim.lr_scheduler import CosineAnnealingLR

model = nn.Linear(10, 5)
optimizer = SGD(model.parameters(), lr=0.1)
scheduler = CosineAnnealingLR(optimizer, T_max=10)

for epoch in range(10):
	scheduler.step()
	current_lr = optimizer.param_groups[0]['lr']
	print(f'Epoch {epoch}: {current_lr}')

В результате на каждой эпохе скорость обучения будет уменьшаться по косинусоиде от 0.1 до 0:

Epoch 0: 0.09765625
Epoch 1: 0.09090909
Epoch 2: 0.0809017
Epoch 3: 0.0690983
Epoch 4: 0.05708685
Epoch 5: 0.04619496
Epoch 6: 0.03732059
Epoch 7: 0.0309017
Epoch 8: 0.02706324
Epoch 9: 0.02598176

Пример

Теперь укажем минимальное значение скорости обучения с помощью параметра eta_min. В этом случае скорость обучения будет колебаться между 0.01 и 0.1:

import torch
import torch.nn as nn
from torch.optim import SGD
from torch.optim.lr_scheduler import CosineAnnealingLR

model = nn.Linear(10, 5)
optimizer = SGD(model.parameters(), lr=0.1)
scheduler = CosineAnnealingLR(optimizer, T_max=10, eta_min=0.01)

for epoch in range(10):
	scheduler.step()
	current_lr = optimizer.param_groups[0]['lr']
	print(f'Epoch {epoch}: {current_lr:.4f}')

Результат выполнения кода:

Epoch 0: 0.0977
Epoch 1: 0.0909
Epoch 2: 0.0809
Epoch 3: 0.0691
Epoch 4: 0.0571
Epoch 5: 0.0462
Epoch 6: 0.0373
Epoch 7: 0.0309
Epoch 8: 0.0271
Epoch 9: 0.0260

Заметьте, что минимальное значение достигается только в конце первого цикла, а затем скорость обучения возрастает снова.

Пример

Рассмотрим использование класса CosineAnnealingLR с оптимизатором Adam в рамках реального процесса обучения. Зафиксируем случайное зерно для воспроизводимости:

import torch
import torch.nn as nn
from torch.optim import Adam
from torch.optim.lr_scheduler import CosineAnnealingLR

torch.manual_seed(0)

model = nn.Sequential(
	nn.Linear(20, 10),
	nn.ReLU(),
	nn.Linear(10, 2)
)
optimizer = Adam(model.parameters(), lr=0.01)
scheduler = CosineAnnealingLR(optimizer, T_max=5, eta_min=0.001)

x = torch.randn(32, 20)
y = torch.randint(0, 2, (32,))

for epoch in range(6):
	scheduler.step()
	outputs = model(x)
	loss = nn.CrossEntropyLoss()(outputs, y)
	optimizer.zero_grad()
	loss.backward()
	optimizer.step()
	
	current_lr = optimizer.param_groups[0]['lr']
	print(f'Epoch {epoch}: lr={current_lr:.4f}, loss={loss.item():.3f}')

Результат выполнения кода:

Epoch 0: lr=0.0096, loss=0.706
Epoch 1: lr=0.0085, loss=0.699
Epoch 2: lr=0.0065, loss=0.691
Epoch 3: lr=0.0035, loss=0.701
Epoch 4: lr=0.0015, loss=0.702
Epoch 5: lr=0.0010, loss=0.684

Смотрите также

класс CosineAnnealingWarmRestarts,
который добавляет "разогрев" и перезапуски косинусного отжига
класс StepLR,
который уменьшает скорость обучения через фиксированные интервалы
класс ReduceLROnPlateau,
который уменьшает скорость обучения при остановке метрики
класс OneCycleLR,
который реализует одноцикловую политику изменения скорости обучения