naver/PUMP

1# Copyright 2022-present NAVER Corp.2# CC BY-NC-SA 4.03# Available only for non-commercial use4 5import os6import torch7import numpy as np8 9 10def mkdir_for(file_path):11    dirname = os.path.split(file_path)[0]12    if dirname: os.makedirs(dirname, exist_ok=True)13    return file_path14 15 16def model_size(model):17    ''' Computes the number of parameters of the model 18    '''19    size = 020    for weights in model.state_dict().values():21        size += np.prod(weights.shape)22    return size23 24 25class cudnn_benchmark:26    " context manager to temporarily disable cudnn benchmark "27    def __init__(self, activate ):28        self.activate = activate29    def __enter__(self):30        self.old_bm = torch.backends.cudnn.benchmark 31        torch.backends.cudnn.benchmark = self.activate32    def __exit__(self, *args):33        torch.backends.cudnn.benchmark = self.old_bm34 35 36def todevice(x, device, non_blocking=False):37    """ Transfer some variables to another device (i.e. GPU, CPU:torch, CPU:numpy).38    x:      array, tensor, or container of such.39    device: pytorch device or 'numpy'40    """41    if isinstance(x, dict):42        return {k:todevice(v, device) for k,v in x.items()}43    44    if isinstance(x, (tuple,list)):45        return type(x)(todevice(e, device) for e in x)46 47    if device == 'numpy':48        if isinstance(x, torch.Tensor):49            x = x.detach().cpu().numpy()50    elif x is not None:51        if isinstance(x, np.ndarray):52            x = torch.from_numpy(x)53        x = x.to(device, non_blocking=non_blocking)54    return x55 56def nparray( x ): return todevice(x, 'numpy')57def cpu( x ): return todevice(x, 'cpu')58def cuda( x ): return todevice(x, 'cuda')59 60 61def image( img, with_trf=False ):62    " convert a torch.Tensor to a numpy image (H, W, 3) "63    def convert_image(img):64        if isinstance(img, torch.Tensor):65            if img.dtype is not torch.uint8:66                img = img * 25567                if img.min() < -10:68                    img = img.clone()69                    for i, (mean, std) in enumerate(zip([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])):70                        img[i] *= std71                        img[i] += 255*mean72                img = img.byte()73            if img.shape[0] <= 3:74                img = img.permute(1,2,0)75        return img76 77    if isinstance(img, tuple):78        if with_trf:79            return nparray(convert_image(img[0])), nparray(img[1])80        else:81            img = img[0]82    return nparray(convert_image(img))83 84 85def image_with_trf( img ):86    return image(img, with_trf=True)87 88class ToTensor:89    " numpy images to float tensors "90    def __call__(self, x):91        assert x.ndim == 4 and x.shape[3] == 392        if isinstance(x, np.ndarray):93            x = torch.from_numpy(x)94        assert x.dtype == torch.uint895        return x.permute(0, 3, 1, 2).float() / 25596