|
|
|
|
|
import math |
|
import numpy as np |
|
from unittest import TestCase |
|
import torch |
|
from fvcore.common.param_scheduler import CosineParamScheduler, MultiStepParamScheduler |
|
from torch import nn |
|
|
|
from detectron2.solver import LRMultiplier, WarmupParamScheduler |
|
|
|
|
|
class TestScheduler(TestCase): |
|
def test_warmup_multistep(self): |
|
p = nn.Parameter(torch.zeros(0)) |
|
opt = torch.optim.SGD([p], lr=5) |
|
|
|
multiplier = WarmupParamScheduler( |
|
MultiStepParamScheduler( |
|
[1, 0.1, 0.01, 0.001], |
|
milestones=[10, 15, 20], |
|
num_updates=30, |
|
), |
|
0.001, |
|
5 / 30, |
|
) |
|
sched = LRMultiplier(opt, multiplier, 30) |
|
|
|
|
|
|
|
|
|
p.sum().backward() |
|
opt.step() |
|
|
|
lrs = [0.005] |
|
for _ in range(30): |
|
sched.step() |
|
lrs.append(opt.param_groups[0]["lr"]) |
|
self.assertTrue(np.allclose(lrs[:5], [0.005, 1.004, 2.003, 3.002, 4.001])) |
|
self.assertTrue(np.allclose(lrs[5:10], 5.0)) |
|
self.assertTrue(np.allclose(lrs[10:15], 0.5)) |
|
self.assertTrue(np.allclose(lrs[15:20], 0.05)) |
|
self.assertTrue(np.allclose(lrs[20:], 0.005)) |
|
|
|
def test_warmup_cosine(self): |
|
p = nn.Parameter(torch.zeros(0)) |
|
opt = torch.optim.SGD([p], lr=5) |
|
multiplier = WarmupParamScheduler( |
|
CosineParamScheduler(1, 0), |
|
0.001, |
|
5 / 30, |
|
) |
|
sched = LRMultiplier(opt, multiplier, 30) |
|
|
|
p.sum().backward() |
|
opt.step() |
|
self.assertEqual(opt.param_groups[0]["lr"], 0.005) |
|
lrs = [0.005] |
|
|
|
for _ in range(30): |
|
sched.step() |
|
lrs.append(opt.param_groups[0]["lr"]) |
|
for idx, lr in enumerate(lrs): |
|
expected_cosine = 2.5 * (1.0 + math.cos(math.pi * idx / 30)) |
|
if idx >= 5: |
|
self.assertAlmostEqual(lr, expected_cosine) |
|
else: |
|
self.assertNotAlmostEqual(lr, expected_cosine) |
|
|
