Source code for mmengine.utils.dl_utils.misc

# Copyright (c) OpenMMLab. All rights reserved.
import pkgutil
from typing import Optional, Tuple, Union

import numpy as np
import torch
import torch.nn as nn

from ..misc import is_tuple_of
from .parrots_wrapper import _BatchNorm, _InstanceNorm

[docs]def is_norm(layer: nn.Module, exclude: Optional[Union[type, Tuple[type]]] = None) -> bool: """Check if a layer is a normalization layer. Args: layer (nn.Module): The layer to be checked. exclude (type, tuple[type], optional): Types to be excluded. Returns: bool: Whether the layer is a norm layer. """ if exclude is not None: if not isinstance(exclude, tuple): exclude = (exclude, ) if not is_tuple_of(exclude, type): raise TypeError( f'"exclude" must be either None or type or a tuple of types, ' f'but got {type(exclude)}: {exclude}') if exclude and isinstance(layer, exclude): return False all_norm_bases = (_BatchNorm, _InstanceNorm, nn.GroupNorm, nn.LayerNorm) return isinstance(layer, all_norm_bases)
[docs]def tensor2imgs(tensor: torch.Tensor, mean: Optional[Tuple[float, float, float]] = None, std: Optional[Tuple[float, float, float]] = None, to_bgr: bool = True): """Convert tensor to 3-channel images or 1-channel gray images. Args: tensor (torch.Tensor): Tensor that contains multiple images, shape ( N, C, H, W). :math:`C` can be either 3 or 1. If C is 3, the format should be RGB. mean (tuple[float], optional): Mean of images. If None, (0, 0, 0) will be used for tensor with 3-channel, while (0, ) for tensor with 1-channel. Defaults to None. std (tuple[float], optional): Standard deviation of images. If None, (1, 1, 1) will be used for tensor with 3-channel, while (1, ) for tensor with 1-channel. Defaults to None. to_bgr (bool): For the tensor with 3 channel, convert its format to BGR. For the tensor with 1 channel, it must be False. Defaults to True. Returns: list[np.ndarray]: A list that contains multiple images. """ assert torch.is_tensor(tensor) and tensor.ndim == 4 channels = tensor.size(1) assert channels in [1, 3] if mean is None: mean = (0, ) * channels if std is None: std = (1, ) * channels assert (channels == len(mean) == len(std) == 3) or \ (channels == len(mean) == len(std) == 1 and not to_bgr) mean = tensor.new_tensor(mean).view(1, -1) std = tensor.new_tensor(std).view(1, -1) tensor = tensor.permute(0, 2, 3, 1) * std + mean imgs = tensor.detach().cpu().numpy() if to_bgr and channels == 3: imgs = imgs[:, :, :, (2, 1, 0)] # RGB2BGR imgs = [np.ascontiguousarray(img) for img in imgs] return imgs
[docs]def has_batch_norm(model: nn.Module) -> bool: """Detect whether model has a BatchNormalization layer. Args: model (nn.Module): training model. Returns: bool: whether model has a BatchNormalization layer """ if isinstance(model, _BatchNorm): return True for m in model.children(): if has_batch_norm(m): return True return False
[docs]def mmcv_full_available() -> bool: """Check whether mmcv-full is installed. Returns: bool: True if mmcv-full is installed else False. """ try: import mmcv # noqa: F401 except ImportError: return False ext_loader = pkgutil.find_loader('mmcv._ext') return ext_loader is not None

© Copyright 2022, mmengine contributors. Revision 0549e1ce.

Built with Sphinx using a theme provided by Read the Docs.
Read the Docs v: v0.8.2
On Read the Docs
Project Home

Free document hosting provided by Read the Docs.