k20hcmus/FishEye8K

1import math2import os3import platform4import subprocess5import time6import warnings7from contextlib import contextmanager8from copy import deepcopy9from pathlib import Path10 11import torch12import torch.distributed as dist13import torch.nn as nn14import torch.nn.functional as F15from torch.nn.parallel import DistributedDataParallel as DDP16 17from utils.general import LOGGER, check_version, colorstr, file_date, git_describe18from utils.lion import Lion19 20LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1))  # https://pytorch.org/docs/stable/elastic/run.html21RANK = int(os.getenv('RANK', -1))22WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))23 24try:25    import thop  # for FLOPs computation26except ImportError:27    thop = None28 29# Suppress PyTorch warnings30warnings.filterwarnings('ignore', message='User provided device_type of \'cuda\', but CUDA is not available. Disabling')31warnings.filterwarnings('ignore', category=UserWarning)32 33 34def smart_inference_mode(torch_1_9=check_version(torch.__version__, '1.9.0')):35    # Applies torch.inference_mode() decorator if torch>=1.9.0 else torch.no_grad() decorator36    def decorate(fn):37        return (torch.inference_mode if torch_1_9 else torch.no_grad)()(fn)38 39    return decorate40 41 42def smartCrossEntropyLoss(label_smoothing=0.0):43    # Returns nn.CrossEntropyLoss with label smoothing enabled for torch>=1.10.044    if check_version(torch.__version__, '1.10.0'):45        return nn.CrossEntropyLoss(label_smoothing=label_smoothing)46    if label_smoothing > 0:47        LOGGER.warning(f'WARNING ⚠️ label smoothing {label_smoothing} requires torch>=1.10.0')48    return nn.CrossEntropyLoss()49 50 51def smart_DDP(model):52    # Model DDP creation with checks53    assert not check_version(torch.__version__, '1.12.0', pinned=True), \54        'torch==1.12.0 torchvision==0.13.0 DDP training is not supported due to a known issue. ' \55        'Please upgrade or downgrade torch to use DDP. See https://github.com/ultralytics/yolov5/issues/8395'56    if check_version(torch.__version__, '1.11.0'):57        return DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK, static_graph=True)58    else:59        return DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)60 61 62def reshape_classifier_output(model, n=1000):63    # Update a TorchVision classification model to class count 'n' if required64    from models.common import Classify65    name, m = list((model.model if hasattr(model, 'model') else model).named_children())[-1]  # last module66    if isinstance(m, Classify):  # YOLOv5 Classify() head67        if m.linear.out_features != n:68            m.linear = nn.Linear(m.linear.in_features, n)69    elif isinstance(m, nn.Linear):  # ResNet, EfficientNet70        if m.out_features != n:71            setattr(model, name, nn.Linear(m.in_features, n))72    elif isinstance(m, nn.Sequential):73        types = [type(x) for x in m]74        if nn.Linear in types:75            i = types.index(nn.Linear)  # nn.Linear index76            if m[i].out_features != n:77                m[i] = nn.Linear(m[i].in_features, n)78        elif nn.Conv2d in types:79            i = types.index(nn.Conv2d)  # nn.Conv2d index80            if m[i].out_channels != n:81                m[i] = nn.Conv2d(m[i].in_channels, n, m[i].kernel_size, m[i].stride, bias=m[i].bias is not None)82 83 84@contextmanager85def torch_distributed_zero_first(local_rank: int):86    # Decorator to make all processes in distributed training wait for each local_master to do something87    if local_rank not in [-1, 0]:88        dist.barrier(device_ids=[local_rank])89    yield90    if local_rank == 0:91        dist.barrier(device_ids=[0])92 93 94def device_count():95    # Returns number of CUDA devices available. Safe version of torch.cuda.device_count(). Supports Linux and Windows96    assert platform.system() in ('Linux', 'Windows'), 'device_count() only supported on Linux or Windows'97    try:98        cmd = 'nvidia-smi -L | wc -l' if platform.system() == 'Linux' else 'nvidia-smi -L | find /c /v ""'  # Windows99        return int(subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1])100    except Exception:101        return 0102 103 104def select_device(device='', batch_size=0, newline=True):105    # device = None or 'cpu' or 0 or '0' or '0,1,2,3'106    s = f'YOLO 🚀 {git_describe() or file_date()} Python-{platform.python_version()} torch-{torch.__version__} '107    device = str(device).strip().lower().replace('cuda:', '').replace('none', '')  # to string, 'cuda:0' to '0'108    cpu = device == 'cpu'109    mps = device == 'mps'  # Apple Metal Performance Shaders (MPS)110    if cpu or mps:111        os.environ['CUDA_VISIBLE_DEVICES'] = '-1'  # force torch.cuda.is_available() = False112    elif device:  # non-cpu device requested113        os.environ['CUDA_VISIBLE_DEVICES'] = device  # set environment variable - must be before assert is_available()114        assert torch.cuda.is_available() and torch.cuda.device_count() >= len(device.replace(',', '')), \115            f"Invalid CUDA '--device {device}' requested, use '--device cpu' or pass valid CUDA device(s)"116 117    if not cpu and not mps and torch.cuda.is_available():  # prefer GPU if available118        devices = device.split(',') if device else '0'  # range(torch.cuda.device_count())  # i.e. 0,1,6,7119        n = len(devices)  # device count120        if n > 1 and batch_size > 0:  # check batch_size is divisible by device_count121            assert batch_size % n == 0, f'batch-size {batch_size} not multiple of GPU count {n}'122        space = ' ' * (len(s) + 1)123        for i, d in enumerate(devices):124            p = torch.cuda.get_device_properties(i)125            s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / (1 << 20):.0f}MiB)\n"  # bytes to MB126        arg = 'cuda:0'127    elif mps and getattr(torch, 'has_mps', False) and torch.backends.mps.is_available():  # prefer MPS if available128        s += 'MPS\n'129        arg = 'mps'130    else:  # revert to CPU131        s += 'CPU\n'132        arg = 'cpu'133 134    if not newline:135        s = s.rstrip()136    LOGGER.info(s)137    return torch.device(arg)138 139 140def time_sync():141    # PyTorch-accurate time142    if torch.cuda.is_available():143        torch.cuda.synchronize()144    return time.time()145 146 147def profile(input, ops, n=10, device=None):148    """ YOLOv5 speed/memory/FLOPs profiler149    Usage:150        input = torch.randn(16, 3, 640, 640)151        m1 = lambda x: x * torch.sigmoid(x)152        m2 = nn.SiLU()153        profile(input, [m1, m2], n=100)  # profile over 100 iterations154    """155    results = []156    if not isinstance(device, torch.device):157        device = select_device(device)158    print(f"{'Params':>12s}{'GFLOPs':>12s}{'GPU_mem (GB)':>14s}{'forward (ms)':>14s}{'backward (ms)':>14s}"159          f"{'input':>24s}{'output':>24s}")160 161    for x in input if isinstance(input, list) else [input]:162        x = x.to(device)163        x.requires_grad = True164        for m in ops if isinstance(ops, list) else [ops]:165            m = m.to(device) if hasattr(m, 'to') else m  # device166            m = m.half() if hasattr(m, 'half') and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m167            tf, tb, t = 0, 0, [0, 0, 0]  # dt forward, backward168            try:169                flops = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2  # GFLOPs170            except Exception:171                flops = 0172 173            try:174                for _ in range(n):175                    t[0] = time_sync()176                    y = m(x)177                    t[1] = time_sync()178                    try:179                        _ = (sum(yi.sum() for yi in y) if isinstance(y, list) else y).sum().backward()180                        t[2] = time_sync()181                    except Exception:  # no backward method182                        # print(e)  # for debug183                        t[2] = float('nan')184                    tf += (t[1] - t[0]) * 1000 / n  # ms per op forward185                    tb += (t[2] - t[1]) * 1000 / n  # ms per op backward186                mem = torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0  # (GB)187                s_in, s_out = (tuple(x.shape) if isinstance(x, torch.Tensor) else 'list' for x in (x, y))  # shapes188                p = sum(x.numel() for x in m.parameters()) if isinstance(m, nn.Module) else 0  # parameters189                print(f'{p:12}{flops:12.4g}{mem:>14.3f}{tf:14.4g}{tb:14.4g}{str(s_in):>24s}{str(s_out):>24s}')190                results.append([p, flops, mem, tf, tb, s_in, s_out])191            except Exception as e:192                print(e)193                results.append(None)194            torch.cuda.empty_cache()195    return results196 197 198def is_parallel(model):199    # Returns True if model is of type DP or DDP200    return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)201 202 203def de_parallel(model):204    # De-parallelize a model: returns single-GPU model if model is of type DP or DDP205    return model.module if is_parallel(model) else model206 207 208def initialize_weights(model):209    for m in model.modules():210        t = type(m)211        if t is nn.Conv2d:212            pass  # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')213        elif t is nn.BatchNorm2d:214            m.eps = 1e-3215            m.momentum = 0.03216        elif t in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:217            m.inplace = True218 219 220def find_modules(model, mclass=nn.Conv2d):221    # Finds layer indices matching module class 'mclass'222    return [i for i, m in enumerate(model.module_list) if isinstance(m, mclass)]223 224 225def sparsity(model):226    # Return global model sparsity227    a, b = 0, 0228    for p in model.parameters():229        a += p.numel()230        b += (p == 0).sum()231    return b / a232 233 234def prune(model, amount=0.3):235    # Prune model to requested global sparsity236    import torch.nn.utils.prune as prune237    for name, m in model.named_modules():238        if isinstance(m, nn.Conv2d):239            prune.l1_unstructured(m, name='weight', amount=amount)  # prune240            prune.remove(m, 'weight')  # make permanent241    LOGGER.info(f'Model pruned to {sparsity(model):.3g} global sparsity')242 243 244def fuse_conv_and_bn(conv, bn):245    # Fuse Conv2d() and BatchNorm2d() layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/246    fusedconv = nn.Conv2d(conv.in_channels,247                          conv.out_channels,248                          kernel_size=conv.kernel_size,249                          stride=conv.stride,250                          padding=conv.padding,251                          dilation=conv.dilation,252                          groups=conv.groups,253                          bias=True).requires_grad_(False).to(conv.weight.device)254 255    # Prepare filters256    w_conv = conv.weight.clone().view(conv.out_channels, -1)257    w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))258    fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.shape))259 260    # Prepare spatial bias261    b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias262    b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))263    fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)264 265    return fusedconv266 267 268def model_info(model, verbose=False, imgsz=640):269    # Model information. img_size may be int or list, i.e. img_size=640 or img_size=[640, 320]270    n_p = sum(x.numel() for x in model.parameters())  # number parameters271    n_g = sum(x.numel() for x in model.parameters() if x.requires_grad)  # number gradients272    if verbose:273        print(f"{'layer':>5} {'name':>40} {'gradient':>9} {'parameters':>12} {'shape':>20} {'mu':>10} {'sigma':>10}")274        for i, (name, p) in enumerate(model.named_parameters()):275            name = name.replace('module_list.', '')276            print('%5g %40s %9s %12g %20s %10.3g %10.3g' %277                  (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))278 279    try:  # FLOPs280        p = next(model.parameters())281        stride = max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32  # max stride282        im = torch.empty((1, p.shape[1], stride, stride), device=p.device)  # input image in BCHW format283        flops = thop.profile(deepcopy(model), inputs=(im,), verbose=False)[0] / 1E9 * 2  # stride GFLOPs284        imgsz = imgsz if isinstance(imgsz, list) else [imgsz, imgsz]  # expand if int/float285        fs = f', {flops * imgsz[0] / stride * imgsz[1] / stride:.1f} GFLOPs'  # 640x640 GFLOPs286    except Exception:287        fs = ''288 289    name = Path(model.yaml_file).stem.replace('yolov5', 'YOLOv5') if hasattr(model, 'yaml_file') else 'Model'290    LOGGER.info(f"{name} summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}")291 292 293def scale_img(img, ratio=1.0, same_shape=False, gs=32):  # img(16,3,256,416)294    # Scales img(bs,3,y,x) by ratio constrained to gs-multiple295    if ratio == 1.0:296        return img297    h, w = img.shape[2:]298    s = (int(h * ratio), int(w * ratio))  # new size299    img = F.interpolate(img, size=s, mode='bilinear', align_corners=False)  # resize300    if not same_shape:  # pad/crop img301        h, w = (math.ceil(x * ratio / gs) * gs for x in (h, w))302    return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447)  # value = imagenet mean303 304 305def copy_attr(a, b, include=(), exclude=()):306    # Copy attributes from b to a, options to only include [...] and to exclude [...]307    for k, v in b.__dict__.items():308        if (len(include) and k not in include) or k.startswith('_') or k in exclude:309            continue310        else:311            setattr(a, k, v)312 313 314def smart_optimizer(model, name='Adam', lr=0.001, momentum=0.9, decay=1e-5):315    # YOLOv5 3-param group optimizer: 0) weights with decay, 1) weights no decay, 2) biases no decay316    g = [], [], []  # optimizer parameter groups317    bn = tuple(v for k, v in nn.__dict__.items() if 'Norm' in k)  # normalization layers, i.e. BatchNorm2d()318    #for v in model.modules():319    #    for p_name, p in v.named_parameters(recurse=0):320    #        if p_name == 'bias':  # bias (no decay)321    #            g[2].append(p)322    #        elif p_name == 'weight' and isinstance(v, bn):  # weight (no decay)323    #            g[1].append(p)324    #        else:325    #            g[0].append(p)  # weight (with decay)326                327    for v in model.modules():328        if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):  # bias (no decay)329            g[2].append(v.bias)330        if isinstance(v, bn):  # weight (no decay)331            g[1].append(v.weight)332        elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):  # weight (with decay)333            g[0].append(v.weight)334            335        if hasattr(v, 'im'):336            if hasattr(v.im, 'implicit'):           337                g[1].append(v.im.implicit)338            else:339                for iv in v.im:340                    g[1].append(iv.implicit)341        if hasattr(v, 'ia'):342            if hasattr(v.ia, 'implicit'):           343                g[1].append(v.ia.implicit)344            else:345                for iv in v.ia:346                    g[1].append(iv.implicit)347                    348        if hasattr(v, 'im2'):349            if hasattr(v.im2, 'implicit'):           350                g[1].append(v.im2.implicit)351            else:352                for iv in v.im2:353                    g[1].append(iv.implicit)354        if hasattr(v, 'ia2'):355            if hasattr(v.ia2, 'implicit'):           356                g[1].append(v.ia2.implicit)357            else:358                for iv in v.ia2:359                    g[1].append(iv.implicit)360                    361        if hasattr(v, 'im3'):362            if hasattr(v.im3, 'implicit'):           363                g[1].append(v.im3.implicit)364            else:365                for iv in v.im3:366                    g[1].append(iv.implicit)367        if hasattr(v, 'ia3'):368            if hasattr(v.ia3, 'implicit'):           369                g[1].append(v.ia3.implicit)370            else:371                for iv in v.ia3:372                    g[1].append(iv.implicit)373                    374        if hasattr(v, 'im4'):375            if hasattr(v.im4, 'implicit'):           376                g[1].append(v.im4.implicit)377            else:378                for iv in v.im4:379                    g[1].append(iv.implicit)380        if hasattr(v, 'ia4'):381            if hasattr(v.ia4, 'implicit'):           382                g[1].append(v.ia4.implicit)383            else:384                for iv in v.ia4:385                    g[1].append(iv.implicit)386                    387        if hasattr(v, 'im5'):388            if hasattr(v.im5, 'implicit'):           389                g[1].append(v.im5.implicit)390            else:391                for iv in v.im5:392                    g[1].append(iv.implicit)393        if hasattr(v, 'ia5'):394            if hasattr(v.ia5, 'implicit'):           395                g[1].append(v.ia5.implicit)396            else:397                for iv in v.ia5:398                    g[1].append(iv.implicit)399                    400        if hasattr(v, 'im6'):401            if hasattr(v.im6, 'implicit'):           402                g[1].append(v.im6.implicit)403            else:404                for iv in v.im6:405                    g[1].append(iv.implicit)406        if hasattr(v, 'ia6'):407            if hasattr(v.ia6, 'implicit'):           408                g[1].append(v.ia6.implicit)409            else:410                for iv in v.ia6:411                    g[1].append(iv.implicit)412                    413        if hasattr(v, 'im7'):414            if hasattr(v.im7, 'implicit'):           415                g[1].append(v.im7.implicit)416            else:417                for iv in v.im7:418                    g[1].append(iv.implicit)419        if hasattr(v, 'ia7'):420            if hasattr(v.ia7, 'implicit'):           421                g[1].append(v.ia7.implicit)422            else:423                for iv in v.ia7:424                    g[1].append(iv.implicit)425 426    if name == 'Adam':427        optimizer = torch.optim.Adam(g[2], lr=lr, betas=(momentum, 0.999))  # adjust beta1 to momentum428    elif name == 'AdamW':429        optimizer = torch.optim.AdamW(g[2], lr=lr, betas=(momentum, 0.999), weight_decay=0.0, amsgrad=True)430    elif name == 'RMSProp':431        optimizer = torch.optim.RMSprop(g[2], lr=lr, momentum=momentum)432    elif name == 'SGD':433        optimizer = torch.optim.SGD(g[2], lr=lr, momentum=momentum, nesterov=True)434    elif name == 'LION':435        optimizer = Lion(g[2], lr=lr, betas=(momentum, 0.99), weight_decay=0.0)436    else:437        raise NotImplementedError(f'Optimizer {name} not implemented.')438 439    optimizer.add_param_group({'params': g[0], 'weight_decay': decay})  # add g0 with weight_decay440    optimizer.add_param_group({'params': g[1], 'weight_decay': 0.0})  # add g1 (BatchNorm2d weights)441    LOGGER.info(f"{colorstr('optimizer:')} {type(optimizer).__name__}(lr={lr}) with parameter groups "442                f"{len(g[1])} weight(decay=0.0), {len(g[0])} weight(decay={decay}), {len(g[2])} bias")443    return optimizer444 445 446def smart_hub_load(repo='ultralytics/yolov5', model='yolov5s', **kwargs):447    # YOLOv5 torch.hub.load() wrapper with smart error/issue handling448    if check_version(torch.__version__, '1.9.1'):449        kwargs['skip_validation'] = True  # validation causes GitHub API rate limit errors450    if check_version(torch.__version__, '1.12.0'):451        kwargs['trust_repo'] = True  # argument required starting in torch 0.12452    try:453        return torch.hub.load(repo, model, **kwargs)454    except Exception:455        return torch.hub.load(repo, model, force_reload=True, **kwargs)456 457 458def smart_resume(ckpt, optimizer, ema=None, weights='yolov5s.pt', epochs=300, resume=True):459    # Resume training from a partially trained checkpoint460    best_fitness = 0.0461    start_epoch = ckpt['epoch'] + 1462    if ckpt['optimizer'] is not None:463        optimizer.load_state_dict(ckpt['optimizer'])  # optimizer464        best_fitness = ckpt['best_fitness']465    if ema and ckpt.get('ema'):466        ema.ema.load_state_dict(ckpt['ema'].float().state_dict())  # EMA467        ema.updates = ckpt['updates']468    if resume:469        assert start_epoch > 0, f'{weights} training to {epochs} epochs is finished, nothing to resume.\n' \470                                f"Start a new training without --resume, i.e. 'python train.py --weights {weights}'"471        LOGGER.info(f'Resuming training from {weights} from epoch {start_epoch} to {epochs} total epochs')472    if epochs < start_epoch:473        LOGGER.info(f"{weights} has been trained for {ckpt['epoch']} epochs. Fine-tuning for {epochs} more epochs.")474        epochs += ckpt['epoch']  # finetune additional epochs475    return best_fitness, start_epoch, epochs476 477 478class EarlyStopping:479    # YOLOv5 simple early stopper480    def __init__(self, patience=30):481        self.best_fitness = 0.0  # i.e. mAP482        self.best_epoch = 0483        self.patience = patience or float('inf')  # epochs to wait after fitness stops improving to stop484        self.possible_stop = False  # possible stop may occur next epoch485 486    def __call__(self, epoch, fitness):487        if fitness >= self.best_fitness:  # >= 0 to allow for early zero-fitness stage of training488            self.best_epoch = epoch489            self.best_fitness = fitness490        delta = epoch - self.best_epoch  # epochs without improvement491        self.possible_stop = delta >= (self.patience - 1)  # possible stop may occur next epoch492        stop = delta >= self.patience  # stop training if patience exceeded493        if stop:494            LOGGER.info(f'Stopping training early as no improvement observed in last {self.patience} epochs. '495                        f'Best results observed at epoch {self.best_epoch}, best model saved as best.pt.\n'496                        f'To update EarlyStopping(patience={self.patience}) pass a new patience value, '497                        f'i.e. `python train.py --patience 300` or use `--patience 0` to disable EarlyStopping.')498        return stop499 500 501class ModelEMA:502    """ Updated Exponential Moving Average (EMA) from https://github.com/rwightman/pytorch-image-models503    Keeps a moving average of everything in the model state_dict (parameters and buffers)504    For EMA details see https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage505    """506 507    def __init__(self, model, decay=0.9999, tau=2000, updates=0):508        # Create EMA509        self.ema = deepcopy(de_parallel(model)).eval()  # FP32 EMA510        self.updates = updates  # number of EMA updates511        self.decay = lambda x: decay * (1 - math.exp(-x / tau))  # decay exponential ramp (to help early epochs)512        for p in self.ema.parameters():513            p.requires_grad_(False)514 515    def update(self, model):516        # Update EMA parameters517        self.updates += 1518        d = self.decay(self.updates)519 520        msd = de_parallel(model).state_dict()  # model state_dict521        for k, v in self.ema.state_dict().items():522            if v.dtype.is_floating_point:  # true for FP16 and FP32523                v *= d524                v += (1 - d) * msd[k].detach()525        # assert v.dtype == msd[k].dtype == torch.float32, f'{k}: EMA {v.dtype} and model {msd[k].dtype} must be FP32'526 527    def update_attr(self, model, include=(), exclude=('process_group', 'reducer')):528        # Update EMA attributes529        copy_attr(self.ema, model, include, exclude)530