Класс StepLR

Класс StepLR является планировщиком скорости обучения в PyTorch. Он уменьшает скорость обучения оптимизатора на заданный коэффициент gamma каждые step_size эпох. Первый параметр принимает оптимизатор, второй параметр задает частоту уменьшения в эпохах.

Синтаксис

torch.optim.lr_scheduler.StepLR(
	optimizer,
	step_size,
	gamma=0.1,
	last_epoch=-1,
	verbose=False
)

Пример

Давайте создадим простую модель и оптимизатор, а затем применим планировщик StepLR для снижения скорости обучения каждые 5 эпох:

import torch
import torch.nn as nn
import torch.optim as optim

model = nn.Linear(10, 1)
optimizer = optim.SGD(model.parameters(), lr=0.1)
scheduler = optim.lr_scheduler.StepLR(
	optimizer, 
	step_size=5, 
	gamma=0.5
)

for epoch in range(12):
	# обучение модели
	optimizer.step()
	scheduler.step()
	current_lr = optimizer.param_groups[0]['lr']
	if epoch % 5 == 0:
		print(f'Эпоха {epoch}: lr = {current_lr}')

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

"Эпоха 0: lr = 0.1"
"Эпоха 5: lr = 0.05"
"Эпоха 10: lr = 0.025"

Пример

Используем StepLR с другими параметрами: уменьшаем скорость обучения на 10% каждые 2 эпохи и выводим сообщения о изменении:

import torch
import torch.nn as nn
import torch.optim as optim

torch.manual_seed(0)
model = nn.Linear(5, 2)
optimizer = optim.Adam(model.parameters(), lr=0.01)
scheduler = optim.lr_scheduler.StepLR(
	optimizer, 
	step_size=2, 
	gamma=0.9,
	verbose=True
)

for epoch in range(6):
	optimizer.step()
	scheduler.step()

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

"Adjusting learning rate of group 0 to 9.0000e-03."
"Adjusting learning rate of group 0 to 8.1000e-03."
"Adjusting learning rate of group 0 to 7.2900e-03."

Пример

Планировщик StepLR можно использовать с любым оптимизатором. Покажем, как работает пошаговое уменьшение с оптимизатором SGD и коэффициентом gamma=0.1 при step_size=3:

import torch
import torch.nn as nn
import torch.optim as optim

model = nn.Linear(2, 2)
optimizer = optim.SGD(model.parameters(), lr=1.0)
scheduler = optim.lr_scheduler.StepLR(
	optimizer, 
	step_size=3, 
	gamma=0.1
)

lrs = []
for epoch in range(10):
	optimizer.step()
	scheduler.step()
	lrs.append(optimizer.param_groups[0]['lr'])

print(f'Скорости обучения: {lrs}')

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

"Скорости обучения: [1.0, 1.0, 1.0, 0.1, 0.1, 0.1, 0.01, 0.01, 0.01, 0.001]"

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

класс MultiStepLR,
который уменьшает LR в заданные моменты времени
класс ExponentialLR,
который уменьшает LR экспоненциально каждую эпоху
класс CosineAnnealingLR,
который изменяет LR по закону косинуса
класс ReduceLROnPlateau,
который уменьшает LR при остановке улучшения метрики