Function divide_param_groups_by_lr_scale

TinyViT/tinyvit_utils.py:79–132

Divide parameters with different lr scale into different groups. Inputs ------ param_groups: a list of dict of torch.nn.Parameter ``` # example: param1.lr_scale = param2.lr_scale = param3.lr_scale = 0.6 param4.lr_scale = param5.lr_scale = param6.lr_scale = 0.3 p

(param_groups)

def divide_param_groups_by_lr_scale(param_groups):
    """
    Divide parameters with different lr scale into different groups.

    Inputs
    ------
    param_groups: a list of dict of torch.nn.Parameter
    ```
    # example:
    param1.lr_scale = param2.lr_scale = param3.lr_scale = 0.6
    param4.lr_scale = param5.lr_scale = param6.lr_scale = 0.3
    param_groups = [{'params': [param1, param2, param4]},
                    {'params': [param3, param5, param6], 'weight_decay': 0.}]

    param_groups = divide_param_groups_by_lr_scale(param_groups)
    ```

    Outputs
    -------
    new_param_groups: a list of dict containing the key `lr_scale`
    ```
    param_groups = [
        {'params': [param1, param2], 'lr_scale': 0.6},
        {'params': [param3], 'weight_decay': 0., 'lr_scale': 0.6}
        {'params': [param4], 'lr_scale': 0.3},
        {'params': [param5, param6], 'weight_decay': 0., 'lr_scale': 0.3}
    ]
    ```
    """
    new_groups = []
    for group in param_groups:
        params = group.pop('params')

        '''
        divide parameters to different groups by lr_scale
        '''
        lr_scale_groups = dict()
        for p in params:
            lr_scale = getattr(p, 'lr_scale', 1.0)

            # create a list if not existed
            if lr_scale not in lr_scale_groups:
                lr_scale_groups[lr_scale] = list()

            # add the parameter with `lr_scale` into the specific group.
            lr_scale_groups[lr_scale].append(p)

        for lr_scale, params in lr_scale_groups.items():
            # copy other parameter information like `weight_decay`
            new_group = copy.copy(group)
            new_group['params'] = params
            new_group['lr_scale'] = lr_scale
            new_groups.append(new_group)
    return new_groups


def set_weight_decay(model):