File size: 4,323 Bytes
6c9ac8f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
# Copyright (c) OpenMMLab. All rights reserved.
from unittest import TestCase
import torch
import torch.nn.functional as F
import torch.optim as optim
from mmengine.optim.scheduler import _ParamScheduler
from mmengine.testing import assert_allclose
from mmdet.engine.schedulers import (QuadraticWarmupLR,
QuadraticWarmupMomentum,
QuadraticWarmupParamScheduler)
class ToyModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = torch.nn.Conv2d(1, 1, 1)
self.conv2 = torch.nn.Conv2d(1, 1, 1)
def forward(self, x):
return self.conv2(F.relu(self.conv1(x)))
class TestQuadraticWarmupScheduler(TestCase):
def setUp(self):
"""Setup the model and optimizer which are used in every test method.
TestCase calls functions in this order: setUp() -> testMethod() ->
tearDown() -> cleanUp()
"""
self.model = ToyModel()
self.optimizer = optim.SGD(
self.model.parameters(), lr=0.05, momentum=0.01, weight_decay=5e-4)
def _test_scheduler_value(self,
schedulers,
targets,
epochs=10,
param_name='lr'):
if isinstance(schedulers, _ParamScheduler):
schedulers = [schedulers]
for epoch in range(epochs):
for param_group, target in zip(self.optimizer.param_groups,
targets):
print(param_group[param_name])
assert_allclose(
target[epoch],
param_group[param_name],
msg='{} is wrong in epoch {}: expected {}, got {}'.format(
param_name, epoch, target[epoch],
param_group[param_name]),
atol=1e-5,
rtol=0)
[scheduler.step() for scheduler in schedulers]
def test_quadratic_warmup_scheduler(self):
with self.assertRaises(ValueError):
QuadraticWarmupParamScheduler(self.optimizer, param_name='lr')
epochs = 10
iters = 5
warmup_factor = [pow((i + 1) / float(iters), 2) for i in range(iters)]
single_targets = [x * 0.05 for x in warmup_factor] + [0.05] * (
epochs - iters)
targets = [single_targets, [x * epochs for x in single_targets]]
scheduler = QuadraticWarmupParamScheduler(
self.optimizer, param_name='lr', end=iters)
self._test_scheduler_value(scheduler, targets, epochs)
def test_quadratic_warmup_scheduler_convert_iterbased(self):
epochs = 10
end = 5
epoch_length = 11
iters = end * epoch_length
warmup_factor = [pow((i + 1) / float(iters), 2) for i in range(iters)]
single_targets = [x * 0.05 for x in warmup_factor] + [0.05] * (
epochs * epoch_length - iters)
targets = [single_targets, [x * epochs for x in single_targets]]
scheduler = QuadraticWarmupParamScheduler.build_iter_from_epoch(
self.optimizer,
param_name='lr',
end=end,
epoch_length=epoch_length)
self._test_scheduler_value(scheduler, targets, epochs * epoch_length)
def test_quadratic_warmup_lr(self):
epochs = 10
iters = 5
warmup_factor = [pow((i + 1) / float(iters), 2) for i in range(iters)]
single_targets = [x * 0.05 for x in warmup_factor] + [0.05] * (
epochs - iters)
targets = [single_targets, [x * epochs for x in single_targets]]
scheduler = QuadraticWarmupLR(self.optimizer, end=iters)
self._test_scheduler_value(scheduler, targets, epochs)
def test_quadratic_warmup_momentum(self):
epochs = 10
iters = 5
warmup_factor = [pow((i + 1) / float(iters), 2) for i in range(iters)]
single_targets = [x * 0.01 for x in warmup_factor] + [0.01] * (
epochs - iters)
targets = [single_targets, [x * epochs for x in single_targets]]
scheduler = QuadraticWarmupMomentum(self.optimizer, end=iters)
self._test_scheduler_value(
scheduler, targets, epochs, param_name='momentum')
|