File size: 523 Bytes
079c32c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
import pytest
import torch
from torch.optim import Adam
from ding.torch_utils.lr_scheduler import cos_lr_scheduler
@pytest.mark.unittest
class TestLRSchedulerHelper:
def test_cos_lr_scheduler(self):
r"""
Overview:
Test the cos lr scheduler.
"""
net = torch.nn.Linear(3, 4)
opt = Adam(net.parameters(), lr=1e-2)
scheduler = cos_lr_scheduler(opt, learning_rate=1e-2, min_lr=6e-5)
scheduler.step(101)
assert opt.param_groups[0]['lr'] == 6e-5
|