train.py 10.8 KB
Newer Older
Kai Chen's avatar
Kai Chen committed
1
import random
2
import re
myownskyW7's avatar
myownskyW7 committed
3
4
from collections import OrderedDict

Kai Chen's avatar
Kai Chen committed
5
import numpy as np
myownskyW7's avatar
myownskyW7 committed
6
import torch
Cao Yuhang's avatar
Cao Yuhang committed
7
import torch.distributed as dist
myownskyW7's avatar
myownskyW7 committed
8
from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
Kai Chen's avatar
Kai Chen committed
9
from mmcv.runner import DistSamplerSeedHook, Runner, obj_from_dict
myownskyW7's avatar
myownskyW7 committed
10

11
from mmdet import datasets
12
13
14
from mmdet.core import (CocoDistEvalmAPHook, CocoDistEvalRecallHook,
                        DistEvalmAPHook, DistOptimizerHook, Fp16OptimizerHook)
from mmdet.datasets import DATASETS, build_dataloader
myownskyW7's avatar
myownskyW7 committed
15
from mmdet.models import RPN
Kai Chen's avatar
Kai Chen committed
16
from mmdet.utils import get_root_logger
myownskyW7's avatar
myownskyW7 committed
17
18


Kai Chen's avatar
Kai Chen committed
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
def set_random_seed(seed, deterministic=False):
    """Set random seed.

    Args:
        seed (int): Seed to be used.
        deterministic (bool): Whether to set the deterministic option for
            CUDNN backend, i.e., set `torch.backends.cudnn.deterministic`
            to True and `torch.backends.cudnn.benchmark` to False.
            Default: False.
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    if deterministic:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


myownskyW7's avatar
myownskyW7 committed
38
39
40
41
42
43
44
45
46
47
48
49
50
51
def parse_losses(losses):
    log_vars = OrderedDict()
    for loss_name, loss_value in losses.items():
        if isinstance(loss_value, torch.Tensor):
            log_vars[loss_name] = loss_value.mean()
        elif isinstance(loss_value, list):
            log_vars[loss_name] = sum(_loss.mean() for _loss in loss_value)
        else:
            raise TypeError(
                '{} is not a tensor or list of tensors'.format(loss_name))

    loss = sum(_value for _key, _value in log_vars.items() if 'loss' in _key)

    log_vars['loss'] = loss
Cao Yuhang's avatar
Cao Yuhang committed
52
53
    for loss_name, loss_value in log_vars.items():
        # reduce loss when distributed training
54
        if dist.is_available() and dist.is_initialized():
Cao Yuhang's avatar
Cao Yuhang committed
55
56
57
            loss_value = loss_value.data.clone()
            dist.all_reduce(loss_value.div_(dist.get_world_size()))
        log_vars[loss_name] = loss_value.item()
myownskyW7's avatar
myownskyW7 committed
58
59
60
61
62

    return loss, log_vars


def batch_processor(model, data, train_mode):
Kai Chen's avatar
Kai Chen committed
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
    """Process a data batch.

    This method is required as an argument of Runner, which defines how to
    process a data batch and obtain proper outputs. The first 3 arguments of
    batch_processor are fixed.

    Args:
        model (nn.Module): A PyTorch model.
        data (dict): The data batch in a dict.
        train_mode (bool): Training mode or not. It may be useless for some
            models.

    Returns:
        dict: A dict containing losses and log vars.
    """
myownskyW7's avatar
myownskyW7 committed
78
79
80
    losses = model(**data)
    loss, log_vars = parse_losses(losses)

81
82
    outputs = dict(
        loss=loss, log_vars=log_vars, num_samples=len(data['img'].data))
myownskyW7's avatar
myownskyW7 committed
83
84
85
86

    return outputs


Kai Chen's avatar
Kai Chen committed
87
88
89
90
91
def train_detector(model,
                   dataset,
                   cfg,
                   distributed=False,
                   validate=False,
92
93
                   timestamp=None):
    logger = get_root_logger(cfg.log_level)
myownskyW7's avatar
myownskyW7 committed
94

Kai Chen's avatar
Kai Chen committed
95
96
    # start training
    if distributed:
97
98
99
100
101
102
103
        _dist_train(
            model,
            dataset,
            cfg,
            validate=validate,
            logger=logger,
            timestamp=timestamp)
myownskyW7's avatar
myownskyW7 committed
104
    else:
105
106
107
108
109
110
111
        _non_dist_train(
            model,
            dataset,
            cfg,
            validate=validate,
            logger=logger,
            timestamp=timestamp)
myownskyW7's avatar
myownskyW7 committed
112

Kai Chen's avatar
Kai Chen committed
113

114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
def build_optimizer(model, optimizer_cfg):
    """Build optimizer from configs.

    Args:
        model (:obj:`nn.Module`): The model with parameters to be optimized.
        optimizer_cfg (dict): The config dict of the optimizer.
            Positional fields are:
                - type: class name of the optimizer.
                - lr: base learning rate.
            Optional fields are:
                - any arguments of the corresponding optimizer type, e.g.,
                  weight_decay, momentum, etc.
                - paramwise_options: a dict with 3 accepted fileds
                  (bias_lr_mult, bias_decay_mult, norm_decay_mult).
                  `bias_lr_mult` and `bias_decay_mult` will be multiplied to
                  the lr and weight decay respectively for all bias parameters
                  (except for the normalization layers), and
                  `norm_decay_mult` will be multiplied to the weight decay
                  for all weight and bias parameters of normalization layers.

    Returns:
        torch.optim.Optimizer: The initialized optimizer.
136
137
138
139
140
141

    Example:
        >>> model = torch.nn.modules.Conv1d(1, 1, 1)
        >>> optimizer_cfg = dict(type='SGD', lr=0.01, momentum=0.9,
        >>>                      weight_decay=0.0001)
        >>> optimizer = build_optimizer(model, optimizer_cfg)
142
143
144
145
146
147
148
149
    """
    if hasattr(model, 'module'):
        model = model.module

    optimizer_cfg = optimizer_cfg.copy()
    paramwise_options = optimizer_cfg.pop('paramwise_options', None)
    # if no paramwise option is specified, just use the global setting
    if paramwise_options is None:
150
151
        return obj_from_dict(optimizer_cfg, torch.optim,
                             dict(params=model.parameters()))
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
    else:
        assert isinstance(paramwise_options, dict)
        # get base lr and weight decay
        base_lr = optimizer_cfg['lr']
        base_wd = optimizer_cfg.get('weight_decay', None)
        # weight_decay must be explicitly specified if mult is specified
        if ('bias_decay_mult' in paramwise_options
                or 'norm_decay_mult' in paramwise_options):
            assert base_wd is not None
        # get param-wise options
        bias_lr_mult = paramwise_options.get('bias_lr_mult', 1.)
        bias_decay_mult = paramwise_options.get('bias_decay_mult', 1.)
        norm_decay_mult = paramwise_options.get('norm_decay_mult', 1.)
        # set param-wise lr and weight decay
        params = []
        for name, param in model.named_parameters():
Cao Yuhang's avatar
Cao Yuhang committed
168
            param_group = {'params': [param]}
169
            if not param.requires_grad:
Cao Yuhang's avatar
Cao Yuhang committed
170
171
172
173
                # FP16 training needs to copy gradient/weight between master
                # weight copy and model weight, it is convenient to keep all
                # parameters here to align with model.parameters()
                params.append(param_group)
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
                continue

            # for norm layers, overwrite the weight decay of weight and bias
            # TODO: obtain the norm layer prefixes dynamically
            if re.search(r'(bn|gn)(\d+)?.(weight|bias)', name):
                if base_wd is not None:
                    param_group['weight_decay'] = base_wd * norm_decay_mult
            # for other layers, overwrite both lr and weight decay of bias
            elif name.endswith('.bias'):
                param_group['lr'] = base_lr * bias_lr_mult
                if base_wd is not None:
                    param_group['weight_decay'] = base_wd * bias_decay_mult
            # otherwise use the global settings

            params.append(param_group)

        optimizer_cls = getattr(torch.optim, optimizer_cfg.pop('type'))
        return optimizer_cls(params, **optimizer_cfg)


194
195
196
197
198
199
def _dist_train(model,
                dataset,
                cfg,
                validate=False,
                logger=None,
                timestamp=None):
myownskyW7's avatar
myownskyW7 committed
200
    # prepare data loaders
201
    dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
myownskyW7's avatar
myownskyW7 committed
202
    data_loaders = [
203
        build_dataloader(
204
205
            ds, cfg.data.imgs_per_gpu, cfg.data.workers_per_gpu, dist=True)
        for ds in dataset
myownskyW7's avatar
myownskyW7 committed
206
207
    ]
    # put model on gpus
Kai Chen's avatar
Kai Chen committed
208
    model = MMDistributedDataParallel(model.cuda())
Cao Yuhang's avatar
Cao Yuhang committed
209

myownskyW7's avatar
myownskyW7 committed
210
    # build runner
211
    optimizer = build_optimizer(model, cfg.optimizer)
212
213
214
215
    runner = Runner(
        model, batch_processor, optimizer, cfg.work_dir, logger=logger)
    # an ugly walkaround to make the .log and .log.json filenames the same
    runner.timestamp = timestamp
Cao Yuhang's avatar
Cao Yuhang committed
216
217
218
219
220
221
222
223
224

    # fp16 setting
    fp16_cfg = cfg.get('fp16', None)
    if fp16_cfg is not None:
        optimizer_config = Fp16OptimizerHook(**cfg.optimizer_config,
                                             **fp16_cfg)
    else:
        optimizer_config = DistOptimizerHook(**cfg.optimizer_config)

myownskyW7's avatar
myownskyW7 committed
225
226
227
    # register hooks
    runner.register_training_hooks(cfg.lr_config, optimizer_config,
                                   cfg.checkpoint_config, cfg.log_config)
Kai Chen's avatar
Kai Chen committed
228
229
230
    runner.register_hook(DistSamplerSeedHook())
    # register eval hooks
    if validate:
231
        val_dataset_cfg = cfg.data.val
232
        eval_cfg = cfg.get('evaluation', {})
Kai Chen's avatar
Kai Chen committed
233
        if isinstance(model.module, RPN):
Kai Chen's avatar
Kai Chen committed
234
            # TODO: implement recall hooks for other datasets
235
236
            runner.register_hook(
                CocoDistEvalRecallHook(val_dataset_cfg, **eval_cfg))
Kai Chen's avatar
Kai Chen committed
237
        else:
238
            dataset_type = DATASETS.get(val_dataset_cfg.type)
239
            if issubclass(dataset_type, datasets.CocoDataset):
240
241
                runner.register_hook(
                    CocoDistEvalmAPHook(val_dataset_cfg, **eval_cfg))
Kai Chen's avatar
Kai Chen committed
242
            else:
243
244
                runner.register_hook(
                    DistEvalmAPHook(val_dataset_cfg, **eval_cfg))
Kai Chen's avatar
Kai Chen committed
245
246
247
248
249
250
251
252

    if cfg.resume_from:
        runner.resume(cfg.resume_from)
    elif cfg.load_from:
        runner.load_checkpoint(cfg.load_from)
    runner.run(data_loaders, cfg.workflow, cfg.total_epochs)


253
254
255
256
257
258
def _non_dist_train(model,
                    dataset,
                    cfg,
                    validate=False,
                    logger=None,
                    timestamp=None):
259
260
261
262
263
    if validate:
        raise NotImplementedError('Built-in validation is not implemented '
                                  'yet in not-distributed training. Use '
                                  'distributed training or test.py and '
                                  '*eval.py scripts instead.')
Kai Chen's avatar
Kai Chen committed
264
    # prepare data loaders
265
    dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
Kai Chen's avatar
Kai Chen committed
266
    data_loaders = [
267
        build_dataloader(
268
            ds,
269
270
271
            cfg.data.imgs_per_gpu,
            cfg.data.workers_per_gpu,
            cfg.gpus,
272
            dist=False) for ds in dataset
Kai Chen's avatar
Kai Chen committed
273
274
275
    ]
    # put model on gpus
    model = MMDataParallel(model, device_ids=range(cfg.gpus)).cuda()
Cao Yuhang's avatar
Cao Yuhang committed
276

Kai Chen's avatar
Kai Chen committed
277
    # build runner
278
    optimizer = build_optimizer(model, cfg.optimizer)
279
280
281
282
    runner = Runner(
        model, batch_processor, optimizer, cfg.work_dir, logger=logger)
    # an ugly walkaround to make the .log and .log.json filenames the same
    runner.timestamp = timestamp
Cao Yuhang's avatar
Cao Yuhang committed
283
284
285
286
287
288
289
290
    # fp16 setting
    fp16_cfg = cfg.get('fp16', None)
    if fp16_cfg is not None:
        optimizer_config = Fp16OptimizerHook(
            **cfg.optimizer_config, **fp16_cfg, distributed=False)
    else:
        optimizer_config = cfg.optimizer_config
    runner.register_training_hooks(cfg.lr_config, optimizer_config,
Kai Chen's avatar
Kai Chen committed
291
                                   cfg.checkpoint_config, cfg.log_config)
myownskyW7's avatar
myownskyW7 committed
292
293
294
295
296

    if cfg.resume_from:
        runner.resume(cfg.resume_from)
    elif cfg.load_from:
        runner.load_checkpoint(cfg.load_from)
myownskyW7's avatar
myownskyW7 committed
297
    runner.run(data_loaders, cfg.workflow, cfg.total_epochs)