mirror of https://github.com/FoundationVision/GLEE
120 lines
3.7 KiB
Python
120 lines
3.7 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
|
|
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, build_lr_scheduler
|
|
|
|
|
|
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)
|
|
# This is an equivalent of:
|
|
# sched = WarmupMultiStepLR(
|
|
# opt, milestones=[10, 15, 20], gamma=0.1, warmup_factor=0.001, warmup_iters=5)
|
|
|
|
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)
|
|
|
|
def test_warmup_cosine_end_value(self):
|
|
from detectron2.config import CfgNode, get_cfg
|
|
|
|
def _test_end_value(cfg_dict):
|
|
cfg = get_cfg()
|
|
cfg.merge_from_other_cfg(CfgNode(cfg_dict))
|
|
|
|
p = nn.Parameter(torch.zeros(0))
|
|
opt = torch.optim.SGD([p], lr=cfg.SOLVER.BASE_LR)
|
|
|
|
scheduler = build_lr_scheduler(cfg, opt)
|
|
|
|
p.sum().backward()
|
|
opt.step()
|
|
self.assertEqual(
|
|
opt.param_groups[0]["lr"], cfg.SOLVER.BASE_LR * cfg.SOLVER.WARMUP_FACTOR
|
|
)
|
|
|
|
lrs = []
|
|
for _ in range(cfg.SOLVER.MAX_ITER):
|
|
scheduler.step()
|
|
lrs.append(opt.param_groups[0]["lr"])
|
|
|
|
self.assertAlmostEqual(lrs[-1], cfg.SOLVER.BASE_LR_END)
|
|
|
|
_test_end_value(
|
|
{
|
|
"SOLVER": {
|
|
"LR_SCHEDULER_NAME": "WarmupCosineLR",
|
|
"MAX_ITER": 100,
|
|
"WARMUP_ITERS": 10,
|
|
"WARMUP_FACTOR": 0.1,
|
|
"BASE_LR": 5.0,
|
|
"BASE_LR_END": 0.0,
|
|
}
|
|
}
|
|
)
|
|
|
|
_test_end_value(
|
|
{
|
|
"SOLVER": {
|
|
"LR_SCHEDULER_NAME": "WarmupCosineLR",
|
|
"MAX_ITER": 100,
|
|
"WARMUP_ITERS": 10,
|
|
"WARMUP_FACTOR": 0.1,
|
|
"BASE_LR": 5.0,
|
|
"BASE_LR_END": 0.5,
|
|
}
|
|
}
|
|
)
|