1
0
mirror of https://github.com/open-mmlab/mmengine.git synced 2025-06-03 20:14:30 +08:00
mmengine/tests/test_hook/test_param_scheduler_hook.py
RangiLyu 9a61b389e7
[Refactor] Add batch_idx to hook input. ()
* [Refactor] Add batch_idx to hook input.

* update
2022-03-29 11:40:38 +08:00

28 lines
780 B
Python

# Copyright (c) OpenMMLab. All rights reserved.
from unittest.mock import Mock
from mmengine.hooks import ParamSchedulerHook
class TestParamSchedulerHook:
def test_after_iter(self):
hook = ParamSchedulerHook()
runner = Mock()
scheduler = Mock()
scheduler.step = Mock()
scheduler.by_epoch = False
runner.param_schedulers = [scheduler]
hook.after_train_iter(runner, 0)
scheduler.step.assert_called()
def test_after_epoch(self):
hook = ParamSchedulerHook()
runner = Mock()
scheduler = Mock()
scheduler.step = Mock()
scheduler.by_epoch = True
runner.param_schedulers = [scheduler]
hook.after_train_epoch(runner)
scheduler.step.assert_called()