Вы можете использовать lr shedular torch.optim.lr_scheduler.StepLR
import torch.optim.lr_scheduler.StepLR
>>> scheduler = StepLR(optimizer, step_size=5, gamma=0.1)
Уменьшает скорость обучения каждой группы параметров по гамме каждые эпохи шага см. Документы здесь Пример из документов
>>> # Assuming optimizer uses lr = 0.05 for all groups
>>> # lr = 0.05 if epoch < 30
>>> # lr = 0.005 if 30 <= epoch < 60
>>> # lr = 0.0005 if 60 <= epoch < 90
>>> # ...
>>> scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
>>> for epoch in range(100):
>>> train(...)
>>> validate(...)
>>> scheduler.step()
Подробнее How to adjust Learning Rate
- torch.optim.lr_scheduler
предоставляет несколько методов для настройки скорости обучения в зависимости от количества эпох.
Обновление:
import torch
import torch.optim as optim
optimizer = optim.SGD([torch.rand((2,2), requires_grad=True)], lr=0.1)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
for epoch in range(1, 21):
scheduler.step()
print('Epoch-{0} lr: {1}'.format(epoch, optimizer.param_groups[0]['lr']))
if epoch % 5 == 0:print()
Epoch-1 lr: 0.1
Epoch-2 lr: 0.1
Epoch-3 lr: 0.1
Epoch-4 lr: 0.1
Epoch-5 lr: 0.1
Epoch-6 lr: 0.010000000000000002
Epoch-7 lr: 0.010000000000000002
Epoch-8 lr: 0.010000000000000002
Epoch-9 lr: 0.010000000000000002
Epoch-10 lr: 0.010000000000000002
Epoch-11 lr: 0.0010000000000000002
Epoch-12 lr: 0.0010000000000000002
Epoch-13 lr: 0.0010000000000000002
Epoch-14 lr: 0.0010000000000000002
Epoch-15 lr: 0.0010000000000000002
Epoch-16 lr: 0.00010000000000000003
Epoch-17 lr: 0.00010000000000000003
Epoch-18 lr: 0.00010000000000000003
Epoch-19 lr: 0.00010000000000000003
Epoch-20 lr: 0.00010000000000000003