Function build_optimizer

classification/optimizer.py:11–87

Build optimizer, set weight decay of normalization to 0 by default.

(config, model)

def build_optimizer(config, model):
    """
    Build optimizer, set weight decay of normalization to 0 by default.
    """
    skip = {}
    skip_keywords = {}
    if hasattr(model, 'no_weight_decay'):
        skip = model.no_weight_decay()
    if hasattr(model, 'no_weight_decay_keywords'):
        skip_keywords = model.no_weight_decay_keywords()

    parameters = set_weight_decay_and_lr(
        model,
        config.TRAIN.WEIGHT_DECAY,
        config.TRAIN.BASE_LR,
        skip,
        skip_keywords,
        lr_layer_decay=config.TRAIN.LR_LAYER_DECAY,
        lr_layer_decay_ratio=config.TRAIN.LR_LAYER_DECAY_RATIO,
        freeze_backbone=config.TRAIN.OPTIMIZER.FREEZE_BACKBONE,
        dcn_lr_mul=config.TRAIN.OPTIMIZER.DCN_LR_MUL,
    )

    opt_lower = config.TRAIN.OPTIMIZER.NAME.lower()
    optimizer = None
    use_zero = config.TRAIN.OPTIMIZER.USE_ZERO
    if use_zero:
        print(f'\nUse Zero!')
        if opt_lower == 'sgd':
            # an ugly implementation
            # this problem is fixed after torch 1.12
            # https://github.com/pytorch/pytorch/issues/71347

            # before 1.12, we could only pass list to zero optimizer, so we first pass parameters[0] with its lr and weight decay,
            # then we add other parameter via parameter group.

            optimizer = ZeroRedundancyOptimizer(
                parameters[0]['params'],
                optimizer_class=optim.SGD,
                momentum=config.TRAIN.OPTIMIZER.MOMENTUM, nesterov=True,
                lr=parameters[0]['lr'], weight_decay=parameters[0]['weight_decay']
            )
            if len(parameters) > 1:
                for param_group in parameters[1:]:
                    optimizer.add_param_group(param_group)
        elif opt_lower == 'adamw':
            optimizer = ZeroRedundancyOptimizer(
                parameters[0]['params'],
                optimizer_class=optim.AdamW,
                eps=config.TRAIN.OPTIMIZER.EPS, betas=config.TRAIN.OPTIMIZER.BETAS,
                lr=parameters[0]['lr'], weight_decay=parameters[0]['weight_decay']
            )
            if len(parameters) > 1:
                for param_group in parameters[1:]:
                    optimizer.add_param_group(param_group)
    else:
        if opt_lower == 'sgd':
            optimizer = optim.SGD(parameters,