mirror of
https://github.com/open-mmlab/mmengine.git
synced 2025-06-03 21:54:44 +08:00
3e3866c1b9
* Support multiple optimizers * minor refinement * improve unit tests * minor fix * Update unit tests for resuming or saving ckpt for multiple optimizers * refine docstring * refine docstring * fix typo * update docstring * refactor the logic to build multiple optimizers * resolve comments * ParamSchedulers spports multiple optimizers * add optimizer_wrapper * fix comment and docstirng * fix unit test * add unit test * refine docstring * RuntimeInfoHook supports printing multi learning rates * resolve comments * add optimizer_wrapper * fix mypy * fix lint * fix OptimizerWrapperDict docstring and add unit test * rename OptimizerWrapper to OptimWrapper, OptimWrapperDict inherit OptimWrapper, and fix as comment * Fix AmpOptimizerWrapper * rename build_optmizer_wrapper to build_optim_wrapper * refine optimizer wrapper * fix AmpOptimWrapper.step, docstring * resolve confict * rename DefaultOptimConstructor * fix as comment * rename clig grad auguments * refactor optim_wrapper config * fix docstring of DefaultOptimWrapperConstructor fix docstring of DefaultOptimWrapperConstructor * add get_lr method to OptimWrapper and OptimWrapperDict * skip some amp unit test * fix unit test * fix get_lr, get_momentum docstring * refactor get_lr, get_momentum, fix as comment * fix error message Co-authored-by: zhouzaida <zhouzaida@163.com>
131 lines
4.9 KiB
Python
131 lines
4.9 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
from unittest import TestCase
|
|
from unittest.mock import Mock
|
|
|
|
import torch.nn as nn
|
|
from torch.optim import SGD
|
|
|
|
from mmengine.hooks import RuntimeInfoHook
|
|
from mmengine.logging import MessageHub
|
|
from mmengine.optim import OptimWrapper, OptimWrapperDict
|
|
|
|
|
|
class TestRuntimeInfoHook(TestCase):
|
|
|
|
def test_before_run(self):
|
|
message_hub = MessageHub.get_instance(
|
|
'runtime_info_hook_test_before_run')
|
|
runner = Mock()
|
|
runner.epoch = 3
|
|
runner.iter = 30
|
|
runner.max_epochs = 4
|
|
runner.max_iters = 40
|
|
runner.message_hub = message_hub
|
|
hook = RuntimeInfoHook()
|
|
hook.before_run(runner)
|
|
self.assertEqual(message_hub.get_info('epoch'), 3)
|
|
self.assertEqual(message_hub.get_info('iter'), 30)
|
|
self.assertEqual(message_hub.get_info('max_epochs'), 4)
|
|
self.assertEqual(message_hub.get_info('max_iters'), 40)
|
|
|
|
def test_before_train(self):
|
|
message_hub = MessageHub.get_instance(
|
|
'runtime_info_hook_test_before_train')
|
|
runner = Mock()
|
|
runner.epoch = 7
|
|
runner.iter = 71
|
|
runner.message_hub = message_hub
|
|
hook = RuntimeInfoHook()
|
|
hook.before_train(runner)
|
|
self.assertEqual(message_hub.get_info('epoch'), 7)
|
|
self.assertEqual(message_hub.get_info('iter'), 71)
|
|
|
|
def test_before_train_epoch(self):
|
|
message_hub = MessageHub.get_instance(
|
|
'runtime_info_hook_test_before_train_epoch')
|
|
runner = Mock()
|
|
runner.epoch = 9
|
|
runner.message_hub = message_hub
|
|
hook = RuntimeInfoHook()
|
|
hook.before_train_epoch(runner)
|
|
self.assertEqual(message_hub.get_info('epoch'), 9)
|
|
|
|
def test_before_train_iter(self):
|
|
model = nn.Linear(1, 1)
|
|
optim1 = SGD(model.parameters(), lr=0.01)
|
|
optim2 = SGD(model.parameters(), lr=0.02)
|
|
optim_wrapper1 = OptimWrapper(optim1)
|
|
optim_wrapper2 = OptimWrapper(optim2)
|
|
optim_wrapper_dict = OptimWrapperDict(
|
|
key1=optim_wrapper1, key2=optim_wrapper2)
|
|
# single optimizer
|
|
message_hub = MessageHub.get_instance(
|
|
'runtime_info_hook_test_before_train_iter')
|
|
runner = Mock()
|
|
runner.iter = 9
|
|
runner.optim_wrapper = optim_wrapper1
|
|
runner.message_hub = message_hub
|
|
hook = RuntimeInfoHook()
|
|
hook.before_train_iter(runner, batch_idx=2, data_batch=None)
|
|
self.assertEqual(message_hub.get_info('iter'), 9)
|
|
self.assertEqual(message_hub.get_scalar('train/lr').current(), 0.01)
|
|
|
|
with self.assertRaisesRegex(AssertionError,
|
|
'runner.optim_wrapper.get_lr()'):
|
|
runner.optim_wrapper = Mock()
|
|
runner.optim_wrapper.get_lr = Mock(return_value='error type')
|
|
hook.before_train_iter(runner, batch_idx=2, data_batch=None)
|
|
|
|
# multiple optimizers
|
|
message_hub = MessageHub.get_instance(
|
|
'runtime_info_hook_test_before_train_iter')
|
|
runner = Mock()
|
|
runner.iter = 9
|
|
optimizer1 = Mock()
|
|
optimizer1.param_groups = [{'lr': 0.01}]
|
|
optimizer2 = Mock()
|
|
optimizer2.param_groups = [{'lr': 0.02}]
|
|
runner.message_hub = message_hub
|
|
runner.optim_wrapper = optim_wrapper_dict
|
|
hook = RuntimeInfoHook()
|
|
hook.before_train_iter(runner, batch_idx=2, data_batch=None)
|
|
self.assertEqual(message_hub.get_info('iter'), 9)
|
|
self.assertEqual(
|
|
message_hub.get_scalar('train/key1.lr').current(), 0.01)
|
|
self.assertEqual(
|
|
message_hub.get_scalar('train/key2.lr').current(), 0.02)
|
|
|
|
def test_after_train_iter(self):
|
|
message_hub = MessageHub.get_instance(
|
|
'runtime_info_hook_test_after_train_iter')
|
|
runner = Mock()
|
|
runner.message_hub = message_hub
|
|
hook = RuntimeInfoHook()
|
|
hook.after_train_iter(
|
|
runner,
|
|
batch_idx=2,
|
|
data_batch=None,
|
|
outputs={'log_vars': {
|
|
'loss_cls': 1.111
|
|
}})
|
|
self.assertEqual(
|
|
message_hub.get_scalar('train/loss_cls').current(), 1.111)
|
|
|
|
def test_after_val_epoch(self):
|
|
message_hub = MessageHub.get_instance(
|
|
'runtime_info_hook_test_after_val_epoch')
|
|
runner = Mock()
|
|
runner.message_hub = message_hub
|
|
hook = RuntimeInfoHook()
|
|
hook.after_val_epoch(runner, metrics={'acc': 0.8})
|
|
self.assertEqual(message_hub.get_scalar('val/acc').current(), 0.8)
|
|
|
|
def test_after_test_epoch(self):
|
|
message_hub = MessageHub.get_instance(
|
|
'runtime_info_hook_test_after_test_epoch')
|
|
runner = Mock()
|
|
runner.message_hub = message_hub
|
|
hook = RuntimeInfoHook()
|
|
hook.after_test_epoch(runner, metrics={'acc': 0.8})
|
|
self.assertEqual(message_hub.get_scalar('test/acc').current(), 0.8)
|