Function construct_optimizer

slowfast/models/optimizer_backbone.py:14–126

Construct a stochastic gradient descent or ADAM optimizer with momentum. Details can be found in: Herbert Robbins, and Sutton Monro. "A stochastic approximation method." and Diederik P.Kingma, and Jimmy Ba. "Adam: A Method for Stochastic Optimization." Args: mod

(model, cfg)

def construct_optimizer(model, cfg):
    """
    Construct a stochastic gradient descent or ADAM optimizer with momentum.
    Details can be found in:
    Herbert Robbins, and Sutton Monro. "A stochastic approximation method."
    and
    Diederik P.Kingma, and Jimmy Ba.
    "Adam: A Method for Stochastic Optimization."

    Args:
        model (model): model to perform stochastic gradient descent
        optimization or ADAM optimization.
        cfg (config): configs of hyper-parameters of SGD or ADAM, includes base
        learning rate,  momentum, weight_decay, dampening, and etc.
    """
    bn_parameters = []
    non_bn_parameters = []
    zero_parameters = []
    head_bn_parameters = []
    head_non_bn_parameters = []
    head_zero_parameters = []
    skip = {}
    if hasattr(model, "no_weight_decay"):
        skip = model.no_weight_decay()

    logger.info(f'LR Ration for backbone is {cfg.SOLVER.BACKBONE_LR_RATIO}')

    total_num = 0
    for name, m in model.named_modules():
        is_bn = isinstance(m, torch.nn.modules.batchnorm._NormBase)
        for p in m.parameters(recurse=False):
            if not p.requires_grad:
                continue
            total_num += 1
            if is_bn:
                if 'backbone.' in name:
                    bn_parameters.append(p)
                else:
                    head_bn_parameters.append(p)
            elif name in skip or (
                (len(p.shape) == 1 or name.endswith(".bias"))
                and cfg.SOLVER.ZERO_WD_1D_PARAM
            ):
                if 'backbone.' in name:
                    zero_parameters.append(p)
                else:
                    head_zero_parameters.append(p)
            else:
                if 'backbone.' in name:
                    non_bn_parameters.append(p)
                else:
                    head_non_bn_parameters.append(p)

    optim_params = [
        {"params": bn_parameters, "weight_decay": cfg.BN.WEIGHT_DECAY, 'lr': cfg.SOLVER.BASE_LR * cfg.SOLVER.BACKBONE_LR_RATIO},
        {"params": non_bn_parameters, "weight_decay": cfg.SOLVER.WEIGHT_DECAY, 'lr': cfg.SOLVER.BASE_LR * cfg.SOLVER.BACKBONE_LR_RATIO},
        {"params": zero_parameters, "weight_decay": 0.0, 'lr': cfg.SOLVER.BASE_LR * cfg.SOLVER.BACKBONE_LR_RATIO},
        {"params": head_bn_parameters, "weight_decay": cfg.BN.WEIGHT_DECAY},