python-wml 3.0.0__py3-none-any.whl

wml/wtorch/__init__.py

@@ -0,0 +1,8 @@
+from .transformer_blocks import TransformerBlock,Attention
+from .conv_module import ConvModule
+from .conv_ws import ConvAWS2d,conv_ws_2d
+from .fc_module import FCModule
+from .summary import *
+from .nn import CHW2HWC,HWC2CHW,LayerNorm,ParallelModule,SumModule,AttentionPool2d
+from .depthwise_separable_conv_module import DepthwiseSeparableConvModule
+from .functional import soft_one_hot

wml/wtorch/bboxes.py

@@ -0,0 +1,104 @@
+import torch
+import torch.nn.functional as F
+
+def cxywh2xy(bboxes):
+    cxy = bboxes[...,:2]
+    hwh = bboxes[...,2:]/2
+    minxy = cxy-hwh
+    maxxy = cxy+hwh
+    return torch.cat([minxy,maxxy],dim=-1)
+
+def xy2cxywh(bboxes):
+    wh = bboxes[...,2:]-bboxes[...,:2]
+    cxy = (bboxes[...,2:]+bboxes[...,:2])/2
+    return torch.cat([cxy,wh],dim=-1)
+
+def distored_boxes(bboxes:torch.Tensor,scale=[0.8,1.2],offset=0.2):
+    bboxes = xy2cxywh(bboxes)
+    cxy,wh = torch.split(bboxes,2,dim=-1)
+    wh_scales = torch.rand(list(wh.shape),dtype=bboxes.dtype)*(scale[1]-scale[0])+scale[0]
+    wh_scales = wh_scales.to(wh.device)
+    wh = wh*wh_scales
+    cxy_offset = torch.rand(list(cxy.shape),dtype=cxy.dtype)*offset
+    cxy_offset = cxy_offset.to(cxy.device)
+    cxy = cxy+cxy_offset
+    bboxes = torch.cat([cxy,wh],axis=-1)
+    bboxes = cxywh2xy(bboxes)
+    bboxes = torch.nn.functional.relu(bboxes)
+    return bboxes
+
+def flip(bboxes,size):
+    '''
+    bboxes:[N,4][xmin,ymin,xmax,ymax]
+    size:[H,W]
+    '''
+    _bboxes = torch.clone(bboxes)
+    _bboxes[...,0] = size[1]-bboxes[...,2]
+    _bboxes[...,2] = size[1]-bboxes[...,0]
+    return _bboxes
+
+def bboxes_ious(bboxesa, bboxesb):
+    '''
+    bboxesa: [N,4] or [1,4] (xmin,ymin,xmax,ymax)
+    bboxesb: [N,4] or [1,4] (xmin,ymin,xmax,ymax)
+    return:
+    [N]
+    '''
+
+    bboxesa = torch.unbind(bboxesa,-1)
+    bboxesb = torch.unbind(bboxesb,-1)
+    int_xmin = torch.maximum(bboxesa[0], bboxesb[0])
+    int_ymin = torch.maximum(bboxesa[1], bboxesb[1])
+    int_xmax = torch.minimum(bboxesa[2], bboxesb[2])
+    int_ymax = torch.minimum(bboxesa[3], bboxesb[3])
+    h = F.relu(int_ymax - int_ymin)
+    w = F.relu(int_xmax - int_xmin)
+    inter_vol = h * w
+    union_vol = -inter_vol \
+                + (bboxesa[2] - bboxesa[0]) * (bboxesa[3] - bboxesa[1]) \
+                + (bboxesb[2] - bboxesb[0]) * (bboxesb[3] - bboxesb[1])
+    union_vol.clamp_(min=1e-8)
+    jaccard = inter_vol/union_vol
+    return jaccard
+
+
+def bboxes_ious_matrix(bboxes0,bboxes1):
+    bboxes0 = torch.unsqueeze(bboxes0,dim=1)
+    bboxes1 = torch.unsqueeze(bboxes1,dim=0)
+
+    x_int_min = torch.maximum(bboxes0[...,0],bboxes1[...,0])
+    x_int_max = torch.minimum(bboxes0[...,2],bboxes1[...,2])
+    y_int_min = torch.maximum(bboxes0[...,1],bboxes1[...,1])
+    y_int_max = torch.minimum(bboxes0[...,3],bboxes1[...,3])
+
+    int_w = x_int_max-x_int_min
+    int_h = y_int_max-y_int_min
+    int_w.clamp_(min=0.0)
+    int_h.clamp_(min=0.0)
+    inter_vol = int_w*int_h
+    areas0 = torch.prod(bboxes0[...,2:]-bboxes0[...,:2],dim=-1)
+    areas1 = torch.prod(bboxes1[...,2:]-bboxes1[...,:2],dim=-1)
+    union_vol = areas0+areas1-inter_vol
+
+    union_vol.clamp_(min=1e-8)
+
+    return inter_vol/union_vol
+
+iou_matrix = bboxes_ious_matrix
+
+def correct_bbox(bboxes,w,h):
+    '''
+    bboxes: [N,4](x0,y0,x1,y1)
+    '''
+    bboxes[:,0:4:2] = torch.clamp(bboxes[:,0:4:2],min=0,max=w)
+    bboxes[:,1:4:2] = torch.clamp(bboxes[:,1:4:2],min=0,max=h)
+
+    return bboxes
+
+def area(bboxes):
+    ws = bboxes[:,2]-bboxes[:,0]
+    hs = bboxes[:,3]-bboxes[:,1]
+    area = ws*hs
+    return area
+
+

wml/wtorch/classes_suppression.py

@@ -0,0 +1,24 @@
+import numpy as np
+
+'''
+根据scores得分,从输入中仅取一个类别
+'''
+def classes_suppression(bboxes,labels,scores,test_nr=1):
+    if len(labels)==0:
+        return bboxes,labels,scores
+    max_scores = -1
+    max_l = -1
+    u_labels = set(labels.tolist())
+    for l in u_labels:
+        m = labels==l
+        idx = np.argsort(-scores[m])[:test_nr]
+        t_scores = np.mean(scores[m][idx])
+        if t_scores>max_scores:
+            max_scores = t_scores
+            max_l = l
+    
+    m = labels==max_l
+
+    return bboxes[m],labels[m],scores[m]
+
+

wml/wtorch/conv_module.py

@@ -0,0 +1,181 @@
+import warnings
+import copy
+import torch.nn as nn
+import wml.wtorch.nn as wnn
+from torch.nn.modules.batchnorm import _BatchNorm
+from torch.nn.modules.instancenorm import _InstanceNorm
+from .nn import get_conv_type
+
+
+class ConvModule(nn.Module):
+    """A conv block that bundles conv/norm/activation layers.
+
+    This block simplifies the usage of convolution layers, which are commonly
+    used with a norm layer (e.g., BatchNorm) and activation layer (e.g., ReLU).
+    It is based upon three build methods: `build_conv_layer()`,
+    `build_norm_layer()` and `build_activation_layer()`.
+
+    Besides, we add some additional features in this module.
+    1. Automatically set `bias` of the conv layer.
+    2. Spectral norm is supported.
+    3. More padding modes are supported. Before PyTorch 1.5, nn.Conv2d only
+    supports zero and circular padding, and we add "reflect" padding mode.
+
+    Args:
+        in_channels (int): Number of channels in the input feature map.
+            Same as that in ``nn._ConvNd``.
+        out_channels (int): Number of channels produced by the convolution.
+            Same as that in ``nn._ConvNd``.
+        kernel_size (int | tuple[int]): Size of the convolving kernel.
+            Same as that in ``nn._ConvNd``.
+        stride (int | tuple[int]): Stride of the convolution.
+            Same as that in ``nn._ConvNd``.
+        padding (int | tuple[int]): Zero-padding added to both sides of
+            the input. Same as that in ``nn._ConvNd``.
+        dilation (int | tuple[int]): Spacing between kernel elements.
+            Same as that in ``nn._ConvNd``.
+        groups (int): Number of blocked connections from input channels to
+            output channels. Same as that in ``nn._ConvNd``.
+        bias (bool | str): If specified as `auto`, it will be decided by the
+            norm_cfg. Bias will be set as True if `norm_cfg` is None, otherwise
+            False. Default: "auto".
+        conv_cfg (dict): Config dict for convolution layer. Default: None,
+            which means using conv2d.
+        norm_cfg (dict): Config dict for normalization layer. Default: None.
+        act_cfg (dict): Config dict for activation layer.
+            Default: dict(type='ReLU').
+        inplace (bool): Whether to use inplace mode for activation.
+            Default: True.
+        with_spectral_norm (bool): Whether use spectral norm in conv module.
+            Default: False.
+        padding_mode (str): If the `padding_mode` has not been supported by
+            current `Conv2d` in PyTorch, we will use our own padding layer
+            instead. Currently, we support ['zeros', 'circular'] with official
+            implementation and ['reflect'] with our own implementation.
+            Default: 'zeros'.
+        order (tuple[str]): The order of conv/norm/activation layers. It is a
+            sequence of "conv", "norm" and "act". Common examples are
+            ("conv", "norm", "act") and ("act", "conv", "norm").
+            Default: ('conv', 'norm', 'act').
+    """
+
+    _abbr_ = 'conv_block'
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 groups=1,
+                 bias='auto',
+                 conv_cfg=None,
+                 norm_cfg=None,
+                 act_cfg=dict(type='ReLU'),
+                 inplace=True,
+                 with_spectral_norm=False,
+                 padding_mode='zeros',
+                 order=('conv', 'norm', 'act')):
+        super(ConvModule, self).__init__()
+        norm_cfg = copy.deepcopy(norm_cfg)
+        assert conv_cfg is None or isinstance(conv_cfg, dict)
+        assert norm_cfg is None or isinstance(norm_cfg, dict)
+        assert act_cfg is None or isinstance(act_cfg, dict)
+        official_padding_mode = ['zeros', 'circular']
+        self.conv_cfg = conv_cfg
+        self.norm_cfg = norm_cfg
+        self.act_cfg = act_cfg
+        self.inplace = inplace
+        self.with_spectral_norm = with_spectral_norm
+        self.with_explicit_padding = padding_mode not in official_padding_mode
+        self.order = order
+        assert isinstance(self.order, tuple) and len(self.order) == 3
+        assert set(order) == set(['conv', 'norm', 'act'])
+
+        self.with_norm = norm_cfg is not None
+        self.with_activation = act_cfg is not None
+        # if the conv layer is before a norm layer, bias is unnecessary.
+        if bias == 'auto':
+            bias = not self.with_norm
+        self.with_bias = bias
+
+        if self.with_explicit_padding:
+            pad_cfg = dict(type=padding_mode)
+            raise RuntimeError("Unsupport with_explicit_padding")
+
+        # reset padding to 0 for conv module
+        conv_padding = 0 if self.with_explicit_padding else padding
+        # build convolution layer
+        self.conv = get_conv_type(conv_cfg)(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=conv_padding,
+            dilation=dilation,
+            groups=groups,
+            bias=bias)
+        # export the attributes of self.conv to a higher level for convenience
+        self.in_channels = self.conv.in_channels
+        self.out_channels = self.conv.out_channels
+        self.kernel_size = self.conv.kernel_size
+        self.stride = self.conv.stride
+        self.padding = padding
+        self.dilation = self.conv.dilation
+        self.transposed = self.conv.transposed
+        self.output_padding = self.conv.output_padding
+        self.groups = self.conv.groups
+
+        if self.with_spectral_norm:
+            self.conv = nn.utils.spectral_norm(self.conv)
+
+        # build normalization layers
+        if self.with_norm:
+            # norm layer is after conv layer
+            if order.index('norm') > order.index('conv'):
+                norm_channels = out_channels
+            else:
+                norm_channels = in_channels
+            norm_type = norm_cfg.pop('type')
+            norm = wnn.get_norm(norm_type,norm_channels,norm_args=norm_cfg)
+            self.norm_name = norm_type
+            self.add_module(self.norm_name, norm)
+            if self.with_bias:
+                if isinstance(norm, (_BatchNorm, _InstanceNorm)):
+                    warnings.warn(
+                        'Unnecessary conv bias before batch/instance norm')
+        else:
+            self.norm_name = None
+
+        # build activation layer
+        if self.with_activation:
+            act_cfg_ = act_cfg.copy()
+            # nn.Tanh has no 'inplace' argument
+            if act_cfg_['type'] not in [
+                    'Tanh', 'PReLU', 'Sigmoid', 'HSigmoid', 'Swish', 'GELU'
+            ]:
+                act_cfg_.setdefault('inplace', inplace)
+            act_type = act_cfg_.pop('type')
+            inplace = act_cfg_.get('inplace',True)
+            self.activate = wnn.get_activation(act_type,inplace)
+
+
+    @property
+    def norm(self):
+        if self.norm_name:
+            return getattr(self, self.norm_name)
+        else:
+            return None
+
+    def forward(self, x, activate=True, norm=True):
+        for layer in self.order:
+            if layer == 'conv':
+                if self.with_explicit_padding:
+                    x = self.padding_layer(x)
+                x = self.conv(x)
+            elif layer == 'norm' and norm and self.with_norm:
+                x = self.norm(x)
+            elif layer == 'act' and activate and self.with_activation:
+                x = self.activate(x)
+        return x

wml/wtorch/conv_ws.py

@@ -0,0 +1,144 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def conv_ws_2d(input,
+               weight,
+               bias=None,
+               stride=1,
+               padding=0,
+               dilation=1,
+               groups=1,
+               eps=1e-5):
+    c_in = weight.size(0)
+    weight_flat = weight.view(c_in, -1)
+    mean = weight_flat.mean(dim=1, keepdim=True).view(c_in, 1, 1, 1)
+    std = weight_flat.std(dim=1, keepdim=True).view(c_in, 1, 1, 1)
+    weight = (weight - mean) / (std + eps)
+    return F.conv2d(input, weight, bias, stride, padding, dilation, groups)
+
+
+class ConvWS2d(nn.Conv2d):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 groups=1,
+                 bias=True,
+                 eps=1e-5):
+        super(ConvWS2d, self).__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            bias=bias)
+        self.eps = eps
+
+    def forward(self, x):
+        return conv_ws_2d(x, self.weight, self.bias, self.stride, self.padding,
+                          self.dilation, self.groups, self.eps)
+
+
+class ConvAWS2d(nn.Conv2d):
+    """AWS (Adaptive Weight Standardization)
+
+    This is a variant of Weight Standardization
+    (https://arxiv.org/pdf/1903.10520.pdf)
+    It is used in DetectoRS to avoid NaN
+    (https://arxiv.org/pdf/2006.02334.pdf)
+
+    Args:
+        in_channels (int): Number of channels in the input image
+        out_channels (int): Number of channels produced by the convolution
+        kernel_size (int or tuple): Size of the conv kernel
+        stride (int or tuple, optional): Stride of the convolution. Default: 1
+        padding (int or tuple, optional): Zero-padding added to both sides of
+            the input. Default: 0
+        dilation (int or tuple, optional): Spacing between kernel elements.
+            Default: 1
+        groups (int, optional): Number of blocked connections from input
+            channels to output channels. Default: 1
+        bias (bool, optional): If set True, adds a learnable bias to the
+            output. Default: True
+    """
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 groups=1,
+                 bias=True):
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            bias=bias)
+        self.register_buffer('weight_gamma',
+                             torch.ones(self.out_channels, 1, 1, 1))
+        self.register_buffer('weight_beta',
+                             torch.zeros(self.out_channels, 1, 1, 1))
+
+    def _get_weight(self, weight):
+        weight_flat = weight.view(weight.size(0), -1)
+        mean = weight_flat.mean(dim=1).view(-1, 1, 1, 1)
+        std = torch.sqrt(weight_flat.var(dim=1) + 1e-5).view(-1, 1, 1, 1)
+        weight = (weight - mean) / std
+        weight = self.weight_gamma * weight + self.weight_beta
+        return weight
+
+    def forward(self, x):
+        weight = self._get_weight(self.weight)
+        return F.conv2d(x, weight, self.bias, self.stride, self.padding,
+                        self.dilation, self.groups)
+
+    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
+                              missing_keys, unexpected_keys, error_msgs):
+        """Override default load function.
+
+        AWS overrides the function _load_from_state_dict to recover
+        weight_gamma and weight_beta if they are missing. If weight_gamma and
+        weight_beta are found in the checkpoint, this function will return
+        after super()._load_from_state_dict. Otherwise, it will compute the
+        mean and std of the pretrained weights and store them in weight_beta
+        and weight_gamma.
+        """
+
+        self.weight_gamma.data.fill_(-1)
+        local_missing_keys = []
+        super()._load_from_state_dict(state_dict, prefix, local_metadata,
+                                      strict, local_missing_keys,
+                                      unexpected_keys, error_msgs)
+        if self.weight_gamma.data.mean() > 0:
+            for k in local_missing_keys:
+                missing_keys.append(k)
+            return
+        weight = self.weight.data
+        weight_flat = weight.view(weight.size(0), -1)
+        mean = weight_flat.mean(dim=1).view(-1, 1, 1, 1)
+        std = torch.sqrt(weight_flat.var(dim=1) + 1e-5).view(-1, 1, 1, 1)
+        self.weight_beta.data.copy_(mean)
+        self.weight_gamma.data.copy_(std)
+        missing_gamma_beta = [
+            k for k in local_missing_keys
+            if k.endswith('weight_gamma') or k.endswith('weight_beta')
+        ]
+        for k in missing_gamma_beta:
+            local_missing_keys.remove(k)
+        for k in local_missing_keys:
+            missing_keys.append(k)

wml/wtorch/data/__init__.py

@@ -0,0 +1,16 @@
+from .sampler import Sampler, SequentialSampler, RandomSampler, SubsetRandomSampler, WeightedRandomSampler, BatchSampler,InfiniteSequentialSampler
+from .dataset import (Dataset, IterableDataset, TensorDataset, ConcatDataset, ChainDataset, BufferedShuffleDataset,
+                      Subset, random_split)
+from .base_data_loader_iter import _BaseDataLoaderIter, _DatasetKind
+from .single_process_data_loader_iter import _SingleProcessDataLoaderIter
+from .multi_processing_data_loader_iter import _MultiProcessingDataLoaderIter
+from .dataset import IterableDataset as IterDataPipe
+from .distributed import DistributedSampler
+from .dataloader import DataLoader, _DatasetKind, get_worker_info
+
+__all__ = ['Sampler', 'SequentialSampler', 'RandomSampler',
+           'SubsetRandomSampler', 'WeightedRandomSampler', 'BatchSampler',
+           'DistributedSampler', 'Dataset', 'IterableDataset', 'TensorDataset',
+           'ConcatDataset', 'ChainDataset', 'BufferedShuffleDataset', 'Subset',
+           'random_split', 'DataLoader', '_DatasetKind', 'get_worker_info',
+           'IterDataPipe']

wml/wtorch/data/_utils/__init__.py

@@ -0,0 +1,45 @@
+r"""Utility classes & functions for data loading. Code in this folder is mostly
+used by ../dataloder.py.
+
+A lot of multiprocessing is used in data loading, which only supports running
+functions defined in global environment (py2 can't serialize static methods).
+Therefore, for code tidiness we put these functions into different files in this
+folder.
+"""
+
+import sys
+import atexit
+
+# old private location of the ExceptionWrapper that some users rely on:
+from torch._utils import ExceptionWrapper
+
+
+IS_WINDOWS = sys.platform == "win32"
+
+
+MP_STATUS_CHECK_INTERVAL = 5.0
+r"""Interval (in seconds) to check status of processes to avoid hanging in
+    multiprocessing data loading. This is mainly used in getting data from
+    another process, in which case we need to periodically check whether the
+    sender is alive to prevent hanging."""
+
+
+python_exit_status = False
+r"""Whether Python is shutting down. This flag is guaranteed to be set before
+the Python core library resources are freed, but Python may already be exiting
+for some time when this is set.
+
+Hook to set this flag is `_set_python_exit_flag`, and is inspired by a similar
+hook in Python 3.7 multiprocessing library:
+https://github.com/python/cpython/blob/d4d60134b29290049e28df54f23493de4f1824b6/Lib/multiprocessing/util.py#L277-L327
+"""
+
+
+def _set_python_exit_flag():
+    global python_exit_status
+    python_exit_status = True
+
+atexit.register(_set_python_exit_flag)
+
+
+from . import worker, signal_handling, pin_memory, collate, fetch