Skip to content

GitLab

"magic_pdf/layout.bak/bbox_sort.py" did not exist on "d5dbed7325d428ce1d8278f19d2030769169bb90"
Switch branch/tag
Find file Normal viewHistoryPermalink
test_lr.py 3.33 KB
Newer Older
Sugon_ldc's avatar
add model TSM  
Sugon_ldc committed
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
110
111
112
113
114
115
116
117
118
119
120
121 
# Copyright (c) OpenMMLab. All rights reserved.
import logging
import os.path as osp
import shutil
import sys
import tempfile
from unittest.mock import MagicMock, call

import torch
import torch.nn as nn
from mmcv.runner import IterTimerHook, PaviLoggerHook, build_runner
from torch.utils.data import DataLoader


def test_tin_lr_updater_hook():
    sys.modules['pavi'] = MagicMock()
    loader = DataLoader(torch.ones((10, 2)))
    runner = _build_demo_runner()

    hook_cfg = dict(type='TINLrUpdaterHook', min_lr=0.1)
    runner.register_hook_from_cfg(hook_cfg)

    hook_cfg = dict(
        type='TINLrUpdaterHook',
        by_epoch=False,
        min_lr=0.1,
        warmup='exp',
        warmup_iters=2,
        warmup_ratio=0.9)
    runner.register_hook_from_cfg(hook_cfg)
    runner.register_hook_from_cfg(dict(type='IterTimerHook'))
    runner.register_hook(IterTimerHook())

    hook_cfg = dict(
        type='TINLrUpdaterHook',
        by_epoch=False,
        min_lr=0.1,
        warmup='constant',
        warmup_iters=2,
        warmup_ratio=0.9)
    runner.register_hook_from_cfg(hook_cfg)
    runner.register_hook_from_cfg(dict(type='IterTimerHook'))
    runner.register_hook(IterTimerHook())

    hook_cfg = dict(
        type='TINLrUpdaterHook',
        by_epoch=False,
        min_lr=0.1,
        warmup='linear',
        warmup_iters=2,
        warmup_ratio=0.9)
    runner.register_hook_from_cfg(hook_cfg)
    runner.register_hook_from_cfg(dict(type='IterTimerHook'))
    runner.register_hook(IterTimerHook())
    # add pavi hook
    hook = PaviLoggerHook(interval=1, add_graph=False, add_last_ckpt=True)
    runner.register_hook(hook)
    runner.run([loader], [('train', 1)])
    shutil.rmtree(runner.work_dir)

    assert hasattr(hook, 'writer')
    calls = [
        call('train', {
            'learning_rate': 0.028544155877284292,
            'momentum': 0.95
        }, 1),
        call('train', {
            'learning_rate': 0.04469266270539641,
            'momentum': 0.95
        }, 6),
        call('train', {
            'learning_rate': 0.09695518130045147,
            'momentum': 0.95
        }, 10)
    ]
    hook.writer.add_scalars.assert_has_calls(calls, any_order=True)


def _build_demo_runner(runner_type='EpochBasedRunner',
                       max_epochs=1,
                       max_iters=None):

    class Model(nn.Module):

        def __init__(self):
            super().__init__()
            self.linear = nn.Linear(2, 1)

        def forward(self, x):
            return self.linear(x)

        def train_step(self, x, optimizer, **kwargs):
            return dict(loss=self(x))

        def val_step(self, x, optimizer, **kwargs):
            return dict(loss=self(x))

    model = Model()

    optimizer = torch.optim.SGD(model.parameters(), lr=0.02, momentum=0.95)

    log_config = dict(
        interval=1, hooks=[
            dict(type='TextLoggerHook'),
        ])

    tmp_dir = tempfile.mkdtemp()
    tmp_dir = osp.join(tmp_dir, '.test_lr_tmp')

    runner = build_runner(
        dict(type=runner_type),
        default_args=dict(
            model=model,
            work_dir=tmp_dir,
            optimizer=optimizer,
            logger=logging.getLogger(),
            max_epochs=max_epochs,
            max_iters=max_iters))
    runner.register_checkpoint_hook(dict(interval=1))
    runner.register_logger_hooks(log_config)
    return runner