Класс CosineAnnealingWarmRestarts

Класс CosineAnnealingWarmRestarts является планировщиком скорости обучения (learning rate scheduler) из модуля torch.optim.lr_scheduler. Он реализует косинусоидальное расписание изменения скорости обучения с периодическими перезапусками. После каждого перезапуска скорость обучения сбрасывается до максимального значения и снова убывает по косинусоиде. Это помогает модели выходить из локальных минимумов и улучшает сходимость.

Планировщик принимает оптимизатор, базовый период перезапуска T_0, количество эпох для увеличения периода T_mult, минимальную скорость обучения eta_min и последнюю эпоху last_epoch. Скорость обучения вычисляется по формуле косинуса и уменьшается от максимального значения до eta_min.

Синтаксис

torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
	optimizer,
	T_0,
	T_mult=1,
	eta_min=0,
	last_epoch=-1,
	verbose=False
)

Основные параметры:

optimizer (Optimizer) - оптимизатор, скорость обучения которого будет изменяться
T_0 (int) - количество эпох до первого перезапуска
T_mult (int) - множитель увеличения периода перезапусков (по умолчанию 1)
eta_min (float) - минимальная скорость обучения (по умолчанию 0)
last_epoch (int) - индекс последней эпохи (по умолчанию -1)
verbose (bool) - выводить ли сообщения о изменении скорости обучения (по умолчанию False)

Пример

Создадим планировщик с базовым периодом T_0 равным 10 эпохам и множителем T_mult равным 2:

import torch
import torch.nn as nn
from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts

model = nn.Linear(10, 1)
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2)

for epoch in range(30):
	optimizer.zero_grad()
	# ... обновление модели ...
	scheduler.step()
	print(f"Epoch {epoch+1}: lr = {optimizer.param_groups[0]['lr']:.6f}")

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

Epoch 1: lr = 0.095106
Epoch 2: lr = 0.080902
Epoch 3: lr = 0.059016
Epoch 4: lr = 0.031251
Epoch 5: lr = 0.000000
Epoch 6: lr = 0.031251
Epoch 7: lr = 0.059016
Epoch 8: lr = 0.080902
Epoch 9: lr = 0.095106
Epoch 10: lr = 0.100000
Epoch 11: lr = 0.097546
Epoch 12: lr = 0.090451
Epoch 13: lr = 0.079390
Epoch 14: lr = 0.065451
...

Пример

Используем планировщик с минимальным значением скорости обучения eta_min равным 0.01:

import torch
import torch.nn as nn
from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts

model = nn.Linear(10, 1)
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
scheduler = CosineAnnealingWarmRestarts(
	optimizer,
	T_0=10,
	T_mult=1,
	eta_min=0.01
)

for epoch in range(25):
	scheduler.step()
	print(f"Epoch {epoch+1}: lr = {optimizer.param_groups[0]['lr']:.6f}")

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

Epoch 1: lr = 0.095106
Epoch 2: lr = 0.080902
Epoch 3: lr = 0.059016
Epoch 4: lr = 0.031251
Epoch 5: lr = 0.010000
Epoch 6: lr = 0.031251
Epoch 7: lr = 0.059016
Epoch 8: lr = 0.080902
Epoch 9: lr = 0.095106
Epoch 10: lr = 0.100000
Epoch 11: lr = 0.095106
Epoch 12: lr = 0.080902
Epoch 13: lr = 0.059016
Epoch 14: lr = 0.031251
Epoch 15: lr = 0.010000
...

Пример

Пример использования планировщика в процессе обучения нейросети на задаче регрессии:

import torch
import torch.nn as nn
from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts

torch.manual_seed(0)

model = nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.5)
scheduler = CosineAnnealingWarmRestarts(
	optimizer,
	T_0=5,
	T_mult=2,
	eta_min=0.001
)

loss_fn = nn.MSELoss()
x = torch.randn(100, 1)
y = x * 2 + 1 + torch.randn(100, 1) * 0.1

for epoch in range(20):
	pred = model(x)
	loss = loss_fn(pred, y)
	optimizer.zero_grad()
	loss.backward()
	optimizer.step()
	scheduler.step()
	cur_lr = optimizer.param_groups[0]['lr']
	print(f"Epoch {epoch+1:2d}: loss = {loss.item():.6f}, lr = {cur_lr:.6f}")

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

Epoch  1: loss = 8.793133, lr = 0.475528
Epoch  2: loss = 5.429678, lr = 0.404509
Epoch  3: loss = 3.354642, lr = 0.295097
Epoch  4: loss = 2.081246, lr = 0.156249
Epoch  5: loss = 1.301461, lr = 0.001000
Epoch  6: loss = 1.087320, lr = 0.156249
Epoch  7: loss = 0.923237, lr = 0.295097
Epoch  8: loss = 0.734662, lr = 0.404509
Epoch  9: loss = 0.574097, lr = 0.475528
Epoch 10: loss = 0.466558, lr = 0.500000
...

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

класс CosineAnnealingLR,
который реализует косинусоидальное расписание без перезапусков
класс StepLR,
который уменьшает скорость обучения на фиксированный коэффициент через заданное число эпох
класс CyclicLR,
который циклически изменяет скорость обучения между двумя границами
класс ReduceLROnPlateau,
который уменьшает скорость обучения при остановке улучшения метрики