Build optimizer, set weight decay of normalization to 0 by default.
(config, model)
| 9 | |
| 10 | |
| 11 | def build_optimizer(config, model): |
| 12 | """ |
| 13 | Build optimizer, set weight decay of normalization to 0 by default. |
| 14 | """ |
| 15 | skip = {} |
| 16 | skip_keywords = {} |
| 17 | if hasattr(model, 'no_weight_decay'): |
| 18 | skip = model.no_weight_decay() |
| 19 | if hasattr(model, 'no_weight_decay_keywords'): |
| 20 | skip_keywords = model.no_weight_decay_keywords() |
| 21 | echo = (config.LOCAL_RANK==0) |
| 22 | parameters = set_weight_decay(model, skip, skip_keywords, echo=echo) |
| 23 | opt_lower = config.TRAIN.OPTIMIZER.NAME.lower() |
| 24 | optimizer = None |
| 25 | if opt_lower == 'sgd': |
| 26 | optimizer = optim.SGD(parameters, momentum=config.TRAIN.OPTIMIZER.MOMENTUM, nesterov=True, |
| 27 | lr=config.TRAIN.BASE_LR, weight_decay=config.TRAIN.WEIGHT_DECAY) |
| 28 | if echo: |
| 29 | print('================================== SGD nest, momentum = {}, wd = {}'.format(config.TRAIN.OPTIMIZER.MOMENTUM, config.TRAIN.WEIGHT_DECAY)) |
| 30 | elif opt_lower == 'adam': |
| 31 | print('adam') |
| 32 | optimizer = optim.Adam(parameters, eps=config.TRAIN.OPTIMIZER.EPS, betas=config.TRAIN.OPTIMIZER.BETAS, |
| 33 | lr=config.TRAIN.BASE_LR, weight_decay=config.TRAIN.WEIGHT_DECAY) |
| 34 | elif opt_lower == 'adamw': |
| 35 | optimizer = optim.AdamW(parameters, eps=config.TRAIN.OPTIMIZER.EPS, betas=config.TRAIN.OPTIMIZER.BETAS, |
| 36 | lr=config.TRAIN.BASE_LR, weight_decay=config.TRAIN.WEIGHT_DECAY) |
| 37 | |
| 38 | return optimizer |
| 39 | |
| 40 | |
| 41 | def set_weight_decay(model, skip_list=(), skip_keywords=(), echo=False): |
no test coverage detected