Метод step

Метод step класса LRScheduler обновляет скорость обучения в соответствии с выбранной стратегией изменения. Обычно этот метод вызывается после optimizer.step для обновления скорости обучения на каждом шаге или после каждой эпохи обучения. Метод автоматически вычисляет новое значение скорости обучения на основе текущего состояния планировщика и применяет его к оптимизатору.

Метод принимает необязательный параметр epoch, который позволяет вручную указать номер эпохи. Это полезно, когда нужно синхронизировать обновление скорости обучения с внешним счётчиком эпох, а не с количеством вызовов метода.

Синтаксис

scheduler.step(epoch=None)

Пример

Создадим простую модель и оптимизатор, а также планировщик, уменьшающий скорость обучения при достижении плато:

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

model = nn.Linear(10, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=2)

print(f"Initial lr: {optimizer.param_groups[0]['lr']}")

for epoch in range(5):
	loss = torch.tensor([5.0 - epoch * 0.5])
	scheduler.step(loss)
	print(f"Epoch {epoch + 1}, lr: {optimizer.param_groups[0]['lr']}")

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

Initial lr: 0.1
Epoch 1, lr: 0.1
Epoch 2, lr: 0.1
Epoch 3, lr: 0.1
Epoch 4, lr: 0.05
Epoch 5, lr: 0.05

Пример

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

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

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

print(f"Initial lr: {optimizer.param_groups[0]['lr']}")

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

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

Initial lr: 0.1
Epoch 1, lr: 0.1
Epoch 2, lr: 0.1
Epoch 3, lr: 0.01
Epoch 4, lr: 0.01
Epoch 5, lr: 0.01
Epoch 6, lr: 0.001

Пример

Используем параметр epoch для ручного управления обновлением скорости обучения:

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

model = nn.Linear(10, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=1.0)
scheduler = ExponentialLR(optimizer, gamma=0.9)

print(f"Initial lr: {optimizer.param_groups[0]['lr']}")

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

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

Initial lr: 1.0
Epoch 1, lr: 0.9
Epoch 2, lr: 0.81
Epoch 3, lr: 0.729
Epoch 4, lr: 0.6561
Epoch 5, lr: 0.59049

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

класс LRScheduler,
базовый класс для всех планировщиков скорости обучения
метод get_last_lr,
возвращает последнее вычисленное значение скорости обучения
метод state_dict,
сохраняет состояние планировщика в словарь
метод load_state_dict,
загружает состояние планировщика из словаря