File size: 1,347 Bytes
ab687e7 |
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 |
from pytorch_caney.lr_scheduler import build_scheduler
import unittest
from unittest.mock import Mock, patch
class TestBuildScheduler(unittest.TestCase):
def setUp(self):
self.config = Mock(
TRAIN=Mock(
EPOCHS=300,
WARMUP_EPOCHS=20,
MIN_LR=1e-6,
WARMUP_LR=1e-7,
LR_SCHEDULER=Mock(
NAME='cosine',
DECAY_EPOCHS=30,
DECAY_RATE=0.1,
MULTISTEPS=[50, 100],
GAMMA=0.1
)
)
)
self.optimizer = Mock()
self.n_iter_per_epoch = 100 # Example value
def test_build_cosine_scheduler(self):
with patch('pytorch_caney.lr_scheduler.CosineLRScheduler') \
as mock_cosine_scheduler:
_ = build_scheduler(self.config,
self.optimizer,
self.n_iter_per_epoch)
mock_cosine_scheduler.assert_called_once_with(
self.optimizer,
t_initial=300 * 100,
cycle_mul=1.,
lr_min=1e-6,
warmup_lr_init=1e-7,
warmup_t=20 * 100,
cycle_limit=1,
t_in_epochs=False
)
if __name__ == '__main__':
unittest.main()
|