util.py

import logging
import os
import numpy as np
import torch
import json
import math
from torch import inf
from decimal import Decimal
if torch.cuda.is_available():
    device = torch.device('cuda')
else:
    device = torch.device('cpu')


def round_num(n, decimal=2):
    n = Decimal(n)
    return round(n.normalize(), decimal)


def numerize(n, decimal=2):
    # https://github.com/davidsa03/numerize/blob/master/numerize/numerize.py
    #60 sufixes
    sufixes = [ "", "K", "M", "B", "T", "Qa", "Qu", "S", "Oc", "No", 
                "D", "Ud", "Dd", "Td", "Qt", "Qi", "Se", "Od", "Nd","V", 
                "Uv", "Dv", "Tv", "Qv", "Qx", "Sx", "Ox", "Nx", "Tn", "Qa",
                "Qu", "S", "Oc", "No", "D", "Ud", "Dd", "Td", "Qt", "Qi",
                "Se", "Od", "Nd", "V", "Uv", "Dv", "Tv", "Qv", "Qx", "Sx",
                "Ox", "Nx", "Tn", "x", "xx", "xxx", "X", "XX", "XXX", "END"] 
    
    sci_expr = [1e0, 1e3, 1e6, 1e9, 1e12, 1e15, 1e18, 1e21, 1e24, 1e27, 
                1e30, 1e33, 1e36, 1e39, 1e42, 1e45, 1e48, 1e51, 1e54, 1e57, 
                1e60, 1e63, 1e66, 1e69, 1e72, 1e75, 1e78, 1e81, 1e84, 1e87, 
                1e90, 1e93, 1e96, 1e99, 1e102, 1e105, 1e108, 1e111, 1e114, 1e117, 
                1e120, 1e123, 1e126, 1e129, 1e132, 1e135, 1e138, 1e141, 1e144, 1e147, 
                1e150, 1e153, 1e156, 1e159, 1e162, 1e165, 1e168, 1e171, 1e174, 1e177]
    minus_buff = n
    n = abs(n)
    if n == 0:
        return str(n)
    for x in range(len(sci_expr)):
        try:
            if n >= sci_expr[x] and n < sci_expr[x+1]:
                sufix = sufixes[x]
                if n >= 1e3:
                    num = str(round_num(n/sci_expr[x], decimal))
                else:
                    num = str(n)
                return num + sufix if minus_buff > 0 else "-" + num + sufix
        except IndexError:
            print("You've reached the end")
    

def update_momentum(model, model_ema, m):
    """Updates parameters of `model_ema` with Exponential Moving Average of `model`
    Momentum encoders are a crucial component fo models such as MoCo or BYOL.
    Examples:
        >>> backbone = resnet18()
        >>> projection_head = MoCoProjectionHead()
        >>> backbone_momentum = copy.deepcopy(moco)
        >>> projection_head_momentum = copy.deepcopy(projection_head)
        >>>
        >>> # update momentum
        >>> update_momentum(moco, moco_momentum, m=0.999)
        >>> update_momentum(projection_head, projection_head_momentum, m=0.999)
    """
    if hasattr(model, 'module'):
        model = model.module
    for model_ema_w, model_w in zip(model_ema.parameters(), model.parameters()):
        model_ema_w.data = model_ema_w.data * m + model_w.data * (1.0 - m)

def adjust_momentum(epoch, config):
    """Adjust moco momentum based on current epoch"""
    if config.momentum_tau_schedule == 'moco':
        m = 1. - 0.5 * (1. + math.cos(math.pi * epoch / config.epochs)) * (1. - config.momentum_tau)
    elif config.momentum_tau_schedule == 'byol':
        m = 1. - (1 - config.momentum_tau) * (math.cos(math.pi * epoch / config.epochs) + 1) / 2
    return m

def adjust_byol_momentum(epoch, config):
    """Adjust moco momentum based on current epoch"""
    m = 1. - 0.5 * (1. + math.cos(math.pi * epoch / config.epochs)) * (1. - config.momentum_tau)
    return m

def exclude_from_wd_and_adaptation(name, bn=True):
        if 'bn' in name and bn:
            return True
        if 'bias' in name:
            return True

def clip_gradients(model, clip):
    norms = []
    for name, p in model.named_parameters():
        if p.grad is not None:
            param_norm = p.grad.data.norm(2)
            norms.append(param_norm.item())
            clip_coef = clip / (param_norm + 1e-6)
            if clip_coef < 1:
                p.grad.data.mul_(clip_coef)
    return norms

def get_params_groups(model):
    regularized = []
    not_regularized = []
    for name, param in model.named_parameters():
        if 'online' in name:
            continue
        if not param.requires_grad:
            continue
        # we do not regularize biases nor Norm parameters
        if name.endswith(".bias") or len(param.shape) == 1:
            not_regularized.append(param)
        else:
            regularized.append(param)
    return [{'params': regularized}, {'params': not_regularized, 'weight_decay': 0.}]


def cancel_gradients_last_layer(epoch, model, freeze_last_layer):
    if epoch >= freeze_last_layer:
        return
    for n, p in model.named_parameters():
        if "last_layer" in n:
            p.grad = None


def get_lars_params(model, weight_decay, bn=True):
    param_groups = [
        {
            'params': [p for name, p in model.named_parameters() if not exclude_from_wd_and_adaptation(name, bn)],
            'weight_decay': weight_decay,
            'layer_adaptation': True,
        },
        {
            'params': [p for name, p in model.named_parameters() if exclude_from_wd_and_adaptation(name, bn)],
            'weight_decay': 0.,
            'layer_adaptation': False,
        },
    ]
    return param_groups


def param_layers_lrd(model, weight_decay=0.05, no_weight_decay_list=[], layer_decay=.75):
    param_group_names = {}
    param_groups = {}
    num_layers = len(list(model.named_parameters())) + 1
    layer_scales = list(layer_decay ** (num_layers - i) for i in range(num_layers + 1))
    for layer_id, (name, p) in enumerate(list(model.named_parameters())):
        if not p.requires_grad:
            continue

        # no decay: all 1D parameters and model specific ones
        this_decay = weight_decay

        if name not in param_group_names:
            this_scale = layer_scales[layer_id]

            param_group_names[name] = {
                "lr_scale": this_scale,
                "weight_decay": this_decay,
                "params": [],
            }
            param_groups[name] = {
                "lr_scale": this_scale,
                "weight_decay": this_decay,
                "params": [],
            }

        param_group_names[name]["params"].append(name)
        param_groups[name]["params"].append(p)

    print("parameter groups: \n%s" % json.dumps(param_group_names, indent=2))
    return list(param_groups.values())


def param_groups_lrd(model, weight_decay=0.05, no_weight_decay_list=[], layer_decay=.75):
    """
    Parameter groups for layer-wise lr decay
    Following BEiT: https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L58
    """
    param_group_names = {}
    param_groups = {}

    num_layers = len(model.blocks) + 1

    layer_scales = list(layer_decay ** (num_layers - i) for i in range(num_layers + 1))

    for n, p in model.named_parameters():
        if not p.requires_grad:
            continue
        if 'predictor' in n:
            continue

        # no decay: all 1D parameters and model specific ones
        if p.ndim == 1 or n in no_weight_decay_list:
            g_decay = "no_decay"
            this_decay = 0.
        else:
            g_decay = "decay"
            this_decay = weight_decay

        layer_id = get_layer_id_for_vit(n, num_layers)
        group_name = "layer_%d_%s" % (layer_id, g_decay)

        if group_name not in param_group_names:
            this_scale = layer_scales[layer_id]

            param_group_names[group_name] = {
                "lr_scale": this_scale,
                "weight_decay": this_decay,
                "params": [],
            }
            param_groups[group_name] = {
                "lr_scale": this_scale,
                "weight_decay": this_decay,
                "params": [],
            }

        param_group_names[group_name]["params"].append(n)
        param_groups[group_name]["params"].append(p)

    # print("parameter groups: \n%s" % json.dumps(param_group_names, indent=2))

    return list(param_groups.values())


def get_layer_id_for_vit(name, num_layers):
    """
    Assign a parameter with its layer id
    Following BEiT: https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L33
    """
    if name in ['cls_token', 'pos_embed']:
        return 0
    elif name.startswith('patch_embed'):
        return 0
    elif name.startswith('blocks'):
        return int(name.split('.')[1]) + 1
    else:
        return num_layers


# --------------------------------------------------------
# Interpolate position embeddings for high-resolution
# References:
# DeiT: https://github.com/facebookresearch/deit
# --------------------------------------------------------
def interpolate_pos_embed(model, checkpoint_model):
    if 'pos_embed' in checkpoint_model:
        pos_embed_checkpoint = checkpoint_model['pos_embed']
        embedding_size = pos_embed_checkpoint.shape[-1]
        num_patches = model.patch_embed.num_patches
        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
        # height (== width) for the checkpoint position embedding
        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
        # height (== width) for the new position embedding
        new_size = int(num_patches ** 0.5)
        # class_token and dist_token are kept unchanged
        if orig_size != new_size:
            print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
            # only the position tokens are interpolated
            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
            pos_tokens = torch.nn.functional.interpolate(
                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
            checkpoint_model['pos_embed'] = new_pos_embed


def has_batchnorms(model):
    bn_types = (torch.nn.BatchNorm1d, torch.nn.BatchNorm2d, 
                torch.nn.BatchNorm3d, torch.nn.SyncBatchNorm)
    for name, module in model.named_modules():
        if isinstance(module, bn_types):
            return True
    return False

def adjust_wd(epoch, configs):
    # Return new WD value used by DINO
    if epoch < configs.warmup_epochs:
        return 0
    q = 0.5 * (1. + math.cos(math.pi * (epoch - configs.warmup_epochs) / (configs.epochs - configs.warmup_epochs)))
    wd = (1 - q) * configs.weight_decay_end + (configs.weight_decay - configs.weight_decay_end) * q
    return wd

def adjust_learning_rate(optimizer, epoch, configs):
    """Decay the learning rate with half-cycle cosine after warmup"""
    if epoch < configs.warmup_epochs:
        lr = configs.lr * epoch / configs.warmup_epochs
    elif 'lr_schedule' in configs:
        if configs.lr_schedule == 'milestone':
            milestones = [int(s*configs.epochs)for s in [0.65, 0.85, 0.95]]
            if epoch < milestones[0]:
                lr = configs.lr
            elif epoch >= milestones[0] and epoch < milestones[1]:
                lr = configs.lr * 0.1
            else:
                lr = configs.lr * 0.01  
        elif configs.lr_schedule == 'linear':
            # lr = np.maximum(configs.lr * np.minimum(configs.epochs / epoch + 1., 1.), 0.)
            lr = configs.lr * (1 - (epoch - configs.warmup_epochs) / (configs.epochs - configs.warmup_epochs))
        elif configs.lr_schedule == 'cosine':
            min_lr = configs.lr * configs.min_lr
            q = 0.5 * (1. + math.cos(math.pi * (epoch - configs.warmup_epochs) / (configs.epochs - configs.warmup_epochs)))
            lr = (1 - q) * min_lr + (configs.lr - min_lr) * q
    else:
        min_lr = configs.lr * configs.min_lr
        q = 0.5 * (1. + math.cos(math.pi * (epoch - configs.warmup_epochs) / (configs.epochs - configs.warmup_epochs)))
        lr = (1 - q) * min_lr + (configs.lr - min_lr) * q

    if hasattr(configs, 'lr_biases'):
        min_lr = configs.lr_biases * configs.min_lr
        q = 0.5 * (1. + math.cos(math.pi * (epoch - configs.warmup_epochs) / (configs.epochs - configs.warmup_epochs)))
        lr_bias = (1 - q) * min_lr + (configs.lr_biases - min_lr) * q
        optimizer.param_groups[0]['lr'] = lr
        optimizer.param_groups[1]['lr'] = lr_bias
    else:
        for param_group in optimizer.param_groups:
            if "lr_scale" in param_group:
                param_group["lr"] = lr * param_group["lr_scale"]
            elif 'fix_lr' in param_group and param_group['fix_lr']:
                param_group["lr"] = configs.lr
            else:
                param_group["lr"] = lr
    return lr


def adjust_learning_rate_with_params(optimizer, epoch, min_lr, lr, warmup, epochs, lr_schedule=None):
    """Decay the learning rate with half-cycle cosine after warmup"""
    if epoch < warmup:
        lr = lr * epoch / warmup
    elif lr_schedule:
        if lr_schedule == 'milestone':
            milestones = [int(s*epochs)for s in [0.65, 0.85, 0.95]]
            if epoch < milestones[0]:
                lr = lr
            elif epoch >= milestones[0] and epoch < milestones[1]:
                lr = lr * 0.1
            else:
                lr = lr * 0.01  
        elif lr_schedule == 'linear':
            lr = lr * (1 - (epoch - warmup) / (epochs - warmup))
    else:
        lr = min_lr + (lr - min_lr) * 0.5 * \
            (1. + math.cos(math.pi * (epoch - warmup) / (epochs - warmup)))
    for param_group in optimizer.param_groups:
        if "lr_scale" in param_group:
            param_group["lr"] = lr * param_group["lr_scale"]
        elif 'fix_lr' in param_group and param_group['fix_lr']:
            print(param_group, lr)
            param_group["lr"] = lr
        else:
            param_group["lr"] = lr
    return lr


def adjust_lid_targets(criterion, epoch, epochs):
    criterion.lid_target = criterion.lid_min + (criterion.lid_max - criterion.lid_min) * 0.5 * \
        (1. + math.cos(math.pi * epoch / epochs))
    return criterion.lid_target


def setup_logger(name, log_file, ddp=False, level=logging.INFO):
    """To setup as many loggers as you want"""
    formatter = logging.Formatter('%(asctime)s %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
    console_handler = logging.StreamHandler()
    console_handler.setFormatter(formatter)
    file_handler = logging.FileHandler(log_file)
    file_handler.setFormatter(formatter)
    logger = logging.getLogger(name)
    logger.setLevel(level)
    logger.addHandler(file_handler)
    if not ddp:
        logger.addHandler(console_handler)
    return logger

def log_display(epoch, global_step, time_elapse, **kwargs):
    display = 'epoch=' + str(epoch) + \
              ' global_step=' + str(numerize(global_step, 1))
    for key, value in kwargs.items():
        if type(value) == str:
            display = ' ' + key + '=' + value
        else:
            if 'acc' in key: 
                display += ' ' + str(key) + '=%.4f' % value
            elif 'loss' in key:
                display += ' ' + str(key) + '=%.4f' % value
            elif 'lr' in key:
                display += ' ' + str(key) + '=%.2e' % value
            else:
                display += ' ' + str(key) + '=%.2f' % value
    display += ' time=%.1fit/s' % (1. / time_elapse)
    return display


def accuracy(output, target, topk=(1,)):
    maxk = max(topk)

    batch_size = target.size(0)
    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))

    res = []
    for k in topk:
        correct_k = correct[:k].float().sum()
        res.append(correct_k.mul_(1/batch_size))
    return res


def mean_cls_acc(output, target):
    acc = [0 for c in range(output.shape[1])]
    _, preds = torch.max(output.data, 1)
    for c in range(output.shape[1]):
        acc[c] = ((preds == target) * (target == c)).sum().float() / max((target == c).sum(), 1)
    return sum(acc) / len(acc)


def count_parameters_in_MB(model):
    return sum(np.prod(v.size()) for name, v in model.named_parameters())/1e6


def build_dirs(path):
    if not os.path.exists(path):
        os.makedirs(path)
    return


class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0
        self.max = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count
        self.max = max(self.max, val)


def get_grad_norm_(parameters, norm_type: float = 2.0) -> torch.Tensor:
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]
    parameters = [p for p in parameters if p.grad is not None]
    norm_type = float(norm_type)
    if len(parameters) == 0:
        return torch.tensor(0.)
    device = parameters[0].grad.device
    if norm_type == inf:
        total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
    else:
        total_norm = torch.norm(torch.stack(
            [torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]), norm_type)
    return total_norm


class NativeScalerWithGradNormCount:
    state_dict_key = "amp_scaler"

    def __init__(self):
        self._scaler = torch.cuda.amp.GradScaler()

    def __call__(self, loss, optimizer, clip_grad=None, parameters=None, create_graph=False, update_grad=True):
        self._scaler.scale(loss).backward(create_graph=create_graph)
        if update_grad:
            if clip_grad is not None:
                assert parameters is not None
                self._scaler.unscale_(optimizer)  # unscale the gradients of optimizer's assigned params in-place
                norm = torch.nn.utils.clip_grad_norm_(parameters, clip_grad)
            else:
                self._scaler.unscale_(optimizer)
                norm = get_grad_norm_(parameters)
            self._scaler.step(optimizer)
            self._scaler.update()
        else:
            norm = None
        return norm

    def state_dict(self):
        return self._scaler.state_dict()

    def load_state_dict(self, state_dict):
        self._scaler.load_state_dict(state_dict)


def unwrap_model(model):
    if hasattr(model, 'module'):
        return model.module
    else:
        return model