python-wml 3.0.0__py3-none-any.whl

wml/wtorch/dropblock/dropblock.py

@@ -0,0 +1,228 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class DropBlock2D(nn.Module):
+    r"""Randomly zeroes 2D spatial blocks of the input tensor.
+
+    As described in the paper
+    `DropBlock: A regularization method for convolutional networks`_ ,
+    dropping whole blocks of feature map allows to remove semantic
+    information as compared to regular dropout.
+
+    Args:
+        drop_prob (float): probability of an element to be dropped.
+        block_size (int): size of the block to drop
+
+    Shape:
+        - Input: `(N, C, H, W)`
+        - Output: `(N, C, H, W)`
+
+    .. _DropBlock: A regularization method for convolutional networks:
+       https://arxiv.org/abs/1810.12890
+
+    """
+
+    def __init__(self, drop_prob, block_size):
+        super(DropBlock2D, self).__init__()
+
+        self.drop_prob = drop_prob
+        self.block_size = block_size
+        self.cur_drop_prob = 0.0
+
+    def forward(self, x):
+        # shape: (bsize, channels, height, width)
+
+        assert x.dim() == 4, \
+            "Expected input with 4 dimensions (bsize, channels, height, width)"
+
+        if not self.training or self.drop_prob <= 0.:
+            self.cur_drop_prob = 0.0
+            return x
+        else:
+            # get gamma value
+            gamma = self._compute_gamma(x)
+
+            # sample mask
+            mask = (torch.rand(x.shape[0], *x.shape[2:]) < gamma).to(x.dtype)
+
+            # place mask on input device
+            mask = mask.to(x.device)
+
+            # compute block mask
+            block_mask = self._compute_block_mask(mask)
+
+            cur_nr = torch.sum(block_mask,dtype=torch.float32)
+            if cur_nr < 1:
+                print(f"ERROR: DropBlock drop_prob={self.drop_prob}, gamma={gamma}, cur_nr={cur_nr}")
+                return x
+
+            scale = (block_mask.numel()/cur_nr).to(x.dtype)
+
+            # apply block mask
+            out = x * block_mask[:, None, :, :]
+
+            # scale output
+            out = out * scale
+
+            self.cur_drop_prob = 1.0-cur_nr/block_mask.numel()
+            return out
+
+    def _compute_block_mask(self, mask):
+        block_mask = F.max_pool2d(input=mask[:, None, :, :],
+                                  kernel_size=(self.block_size, self.block_size),
+                                  stride=(1, 1),
+                                  padding=self.block_size // 2)
+
+        if self.block_size % 2 == 0:
+            block_mask = block_mask[:, :, :-1, :-1]
+
+        block_mask = 1 - block_mask.squeeze(1)
+
+        return block_mask
+
+    def _compute_gamma(self, x):
+        return self.drop_prob / (self.block_size ** 2)
+
+
+class WDropBlock2D(nn.Module):
+    r"""
+    与DropBlock2D不同的在于：DropBlock2d会将一个空间位置整个通道都设置为0，或者都设置为1,
+    WDropBlock2D会将channel中的各元素单独drop
+    """
+
+    def __init__(self, drop_prob, block_size):
+        super().__init__()
+
+        self.drop_prob = drop_prob
+        self.block_size = block_size
+        self.cur_drop_prob = 0.0
+
+    def forward(self, x):
+        # shape: (bsize, channels, height, width)
+
+        assert x.dim() == 4, \
+            "Expected input with 4 dimensions (bsize, channels, height, width)"
+
+        if not self.training or self.drop_prob <= 0.:
+            self.cur_drop_prob = 0.0
+            return x
+        else:
+            # get gamma value
+            gamma = self._compute_gamma(x)
+
+            # sample mask
+            mask = (torch.rand(*x.shape) < gamma).to(x.dtype)
+
+            # place mask on input device
+            mask = mask.to(x.device)
+
+            # compute block mask
+            block_mask = self._compute_block_mask(mask)
+
+            cur_nr = torch.sum(block_mask,dtype=torch.float32)
+            if cur_nr < 1:
+                print(f"ERROR: DropBlock drop_prob={self.drop_prob}, gamma={gamma}, cur_nr={cur_nr}")
+                return x
+
+            scale = (block_mask.numel()/cur_nr).to(x.dtype)
+
+            # apply block mask
+            out = x * block_mask
+
+            # scale output
+            out = out * scale
+
+            self.cur_drop_prob = 1.0-cur_nr/block_mask.numel()
+            return out
+
+    def _compute_block_mask(self, mask):
+        block_mask = F.max_pool2d(input=mask,
+                                  kernel_size=(self.block_size, self.block_size),
+                                  stride=(1, 1),
+                                  padding=self.block_size // 2)
+
+        if self.block_size % 2 == 0:
+            block_mask = block_mask[:, :, :-1, :-1]
+
+        block_mask = 1 - block_mask
+
+        return block_mask
+
+    def _compute_gamma(self, x):
+        return self.drop_prob / (self.block_size ** 2)
+
+
+class DropBlock3D(DropBlock2D):
+    r"""Randomly zeroes 3D spatial blocks of the input tensor.
+
+    An extension to the concept described in the paper
+    `DropBlock: A regularization method for convolutional networks`_ ,
+    dropping whole blocks of feature map allows to remove semantic
+    information as compared to regular dropout.
+
+    Args:
+        drop_prob (float): probability of an element to be dropped.
+        block_size (int): size of the block to drop
+
+    Shape:
+        - Input: `(N, C, D, H, W)`
+        - Output: `(N, C, D, H, W)`
+
+    .. _DropBlock: A regularization method for convolutional networks:
+       https://arxiv.org/abs/1810.12890
+
+    """
+
+    def __init__(self, drop_prob, block_size):
+        super(DropBlock3D, self).__init__(drop_prob, block_size)
+
+    def forward(self, x):
+        # shape: (bsize, channels, depth, height, width)
+
+        assert x.dim() == 5, \
+            "Expected input with 5 dimensions (bsize, channels, depth, height, width)"
+
+        if not self.training or self.drop_prob == 0.:
+            return x
+        else:
+            # get gamma value
+            gamma = self._compute_gamma(x)
+
+            # sample mask
+            mask = (torch.rand(x.shape[0], *x.shape[2:]) < gamma).float()
+
+            # place mask on input device
+            mask = mask.to(x.device)
+
+            # compute block mask
+            block_mask = self._compute_block_mask(mask)
+
+            # apply block mask
+            out = x * block_mask[:, None, :, :, :]
+
+            cur_nr = block_mask.sum()
+            if cur_nr < 1:
+                print(f"ERROR: DropBlock drop_prob={self.drop_prob}, gamma={gamma}, cur_nr={cur_nr}")
+                return x
+            # scale output
+            out = out * block_mask.numel() / cur_nr 
+
+            return out
+
+    def _compute_block_mask(self, mask):
+        block_mask = F.max_pool3d(input=mask[:, None, :, :, :],
+                                  kernel_size=(self.block_size, self.block_size, self.block_size),
+                                  stride=(1, 1, 1),
+                                  padding=self.block_size // 2)
+
+        if self.block_size % 2 == 0:
+            block_mask = block_mask[:, :, :-1, :-1, :-1]
+
+        block_mask = 1 - block_mask.squeeze(1)
+
+        return block_mask
+
+    def _compute_gamma(self, x):
+        return self.drop_prob / (self.block_size ** 3)

wml/wtorch/dropblock/dropout.py

@@ -0,0 +1,40 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class WDropout(nn.Module):
+    r"""Randomly zeroes 2D spatial blocks of the input tensor.
+
+    As described in the paper
+    `DropBlock: A regularization method for convolutional networks`_ ,
+    dropping whole blocks of feature map allows to remove semantic
+    information as compared to regular dropout.
+
+    Args:
+        drop_prob (float): probability of an element to be dropped.
+        block_size (int): size of the block to drop
+
+    Shape:
+        - Input: `(N, C, H, W)`
+        - Output: `(N, C, H, W)`
+
+    .. _DropBlock: A regularization method for convolutional networks:
+       https://arxiv.org/abs/1810.12890
+
+    """
+
+    def __init__(self, drop_prob,inplace=False):
+        super().__init__()
+
+        self.drop_prob = drop_prob
+        self.inplace = inplace
+        self.cur_drop_prob = 0.0
+
+    def forward(self, x):
+        if not self.training or self.drop_prob <= 0.:
+            self.cur_drop_prob = 0.0
+            return x
+        else:
+            self.cur_drop_prob = self.drop_prob
+            return F.dropout(x, self.drop_prob, self.training, self.inplace)

wml/wtorch/dropblock/scheduler.py

@@ -0,0 +1,48 @@
+import numpy as np
+from torch import nn
+
+
+class LinearScheduler(nn.Module):
+    '''
+    在指定的步数范围内将dropblock的prob设置为相应的值
+    如果start_value>0,stop_value>0则设置为对应的插值
+    否则设置为dropblock的默认值
+    '''
+    def __init__(self, dropblock, start_value=-1.0, stop_value=-1.0, begin_step=0,end_step=-1,enable_prob=None):
+        super(LinearScheduler, self).__init__()
+        self.dropblock = dropblock
+        self.i = 0
+        if start_value>=0.0 and stop_value>=0.0 and end_step>0:
+            nr_steps = end_step-begin_step
+            assert nr_steps>0,f"ERROR begin step and end {begin_step} {end_step}"
+            self.drop_values = np.linspace(start=start_value, stop=stop_value, num=int(nr_steps))
+            self.drop_prob = None
+        elif start_value<0 and stop_value<0:
+            self.drop_values = None
+            self.drop_prob = self.dropblock.drop_prob
+            if end_step<0:
+                end_step = 1e9
+        self.begin_step = begin_step
+        self.end_step = end_step
+        self.enable_prob = enable_prob
+
+    def forward(self, x):
+        if self.enable_prob is not None and self.enable_prob > 0:
+            if np.random.rand() > self.enable_prob:
+                self.dropblock.cur_drop_prob = 0.0
+                return x
+        return self.dropblock(x)
+
+    def step(self,step=None):
+        if step is not None:
+            self.i = step
+        else:
+            self.i += 1
+        if self.i>=self.begin_step and self.i<=self.end_step:
+            if self.drop_prob is not None:
+                self.dropblock.drop_prob = self.drop_prob
+            elif self.drop_values is not None and self.i-self.begin_step < len(self.drop_values):
+                self.dropblock.drop_prob = self.drop_values[self.i-self.begin_step]
+        elif self.drop_prob is not None:
+            self.dropblock.drop_prob = 0.0
+

wml/wtorch/ema.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
+import math
+from copy import deepcopy
+
+import torch
+import torch.nn as nn
+
+__all__ = ["ModelEMA", "is_parallel"]
+
+
+def is_parallel(model):
+    """check if model is in parallel mode."""
+    parallel_type = (
+        nn.parallel.DataParallel,
+        nn.parallel.DistributedDataParallel,
+    )
+    return isinstance(model, parallel_type)
+
+
+class ModelEMA:
+    """
+    Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models
+    Keep a moving average of everything in the model state_dict (parameters and buffers).
+    This is intended to allow functionality like
+    https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
+    A smoothed version of the weights is necessary for some training schemes to perform well.
+    This class is sensitive where it is initialized in the sequence of model init,
+    GPU assignment and distributed training wrappers.
+    """
+
+    def __init__(self, model, decay=0.9999, updates=0):
+        """
+        Args:
+            model (nn.Module): model to apply EMA.
+            decay (float): ema decay reate.
+            updates (int): counter of EMA updates.
+        """
+        # Create EMA(FP32)
+        self.ema = deepcopy(model.module if is_parallel(model) else model).eval()
+        self.updates = updates
+        # decay exponential ramp (to help early epochs)
+        self.decay = lambda x: decay * (1 - math.exp(-x / 2000))
+        for p in self.ema.parameters():
+            p.requires_grad_(False)
+
+    def update(self, model):
+        # Update EMA parameters
+        #print(f"EMA Update.")
+        with torch.no_grad():
+            self.updates += 1
+            d = self.decay(self.updates) #old persent
+
+            msd = (
+                model.module.state_dict() if is_parallel(model) else model.state_dict()
+            )  # model state_dict
+            for k, v in self.ema.state_dict().items():
+                if v.dtype.is_floating_point:
+                    v *= d
+                    v += (1.0 - d) * msd[k].detach()

wml/wtorch/fc_module.py

@@ -0,0 +1,73 @@
+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
+
+
+
+class FCModule(nn.Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 bias='auto',
+                 norm_cfg=None,
+                 act_cfg=dict(type='ReLU'),
+                 inplace=True,
+                 ):
+        super().__init__()
+        norm_cfg = copy.deepcopy(norm_cfg)
+        assert norm_cfg is None or isinstance(norm_cfg, dict)
+        assert act_cfg is None or isinstance(act_cfg, dict)
+        self.norm_cfg = norm_cfg
+        self.act_cfg = act_cfg
+
+        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
+
+        modules = []
+        # build convolution layer
+        fc = nn.Linear(
+            in_channels,
+            out_channels,
+            bias=bias)
+        # export the attributes of self.conv to a higher level for convenience
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        modules.append(fc)
+
+        # build normalization layers
+        if self.with_norm:
+            # norm layer is after conv layer
+            norm_type = norm_cfg.pop('type')
+            norm = wnn.get_norm1d(norm_type,
+                                  self.out_channels,
+                                  norm_args=norm_cfg)
+            if self.with_bias:
+                if isinstance(norm, (_BatchNorm, _InstanceNorm)):
+                    warnings.warn(
+                        'Unnecessary conv bias before batch/instance norm')
+            modules.append(norm)
+
+        # 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_.pop('inplace')
+            activate = wnn.get_activation(act_type,inplace)
+            modules.append(activate)
+        self.fc = nn.Sequential(*modules)
+
+
+    def forward(self, x):
+        return self.fc.forward(x)

wml/wtorch/functional.py

@@ -0,0 +1,34 @@
+import torch
+
+def soft_one_hot(labels,confidence,smooth_cfg,num_classes):
+    '''
+    labels: any shape in, value in [0,num_classes)
+    smooth_cfg:Example:
+    {1:[3,4], #标签1向3，4平滑
+     3:[1,4],
+     5:{2:0.04,3:0.06}, #标签5向2,3平滑，2得到0.04的置信度，3得到0.06的置信度
+    }
+    '''
+    old_shape = labels.shape
+    labels = torch.reshape(labels,[-1])
+    targets = torch.zeros([labels.shape[0],num_classes],dtype=torch.float32)
+    smooth_labels = set(list(smooth_cfg.keys()))
+    h_labels = labels.cpu().numpy()
+    targets[torch.arange(targets.shape[0]),h_labels] = 1.0
+    for i,l in enumerate(h_labels):
+        if l in smooth_labels:
+            d = smooth_cfg[l]
+            if isinstance(d,dict):
+                for k,v in d.items():
+                    targets[i,k] = v
+            else:
+                smooth_v = (1-confidence)/len(d)
+                for k in d:
+                    targets[i,k] = smooth_v
+            targets[i,l] = confidence
+
+    targets = targets.to(labels.device)
+    ret_shape = list(old_shape)+[num_classes]
+    targets = torch.reshape(targets,ret_shape)
+
+    return targets

wml/wtorch/iter_dataset.py

@@ -0,0 +1,26 @@
+import os
+class IterdatasetWrapper:
+    def __init__(self,dataset,max_iter=100000):
+        self.dataset = dataset
+        self.max_iter = max_iter
+        self.iter = iter(self.dataset)
+
+    def __len__(self):
+        return self.max_iter
+
+    def __getitem__(self, item):
+        try:
+            data = next(self.iter)
+        except StopIteration as e:
+            self.iter = iter(self.dataset)
+            try:
+                data = next(self.iter)
+            except Exception as e:
+                print('ERROR: IterdatasetWrapper:',e)
+                raise e
+        except Exception as e:
+            print('ERROR: IterdatasetWrapper:',e)
+            raise e
+        return data
+
+

wml/wtorch/loss.py

@@ -0,0 +1,69 @@
+import torch
+import torch.nn.functional as F
+
+def focal_loss_for_heat_map(labels,logits,pos_threshold=0.99,alpha=2,beta=4,sum=True):
+    '''
+    focal loss for heat map, for example CenterNet2's heat map loss
+    '''
+    logits = logits.to(torch.float32)
+    zeros = torch.zeros_like(labels)
+    ones = torch.ones_like(labels)
+    num_pos = torch.sum(torch.where(torch.greater_equal(labels, pos_threshold), ones, zeros))
+
+    probs = F.sigmoid(logits)
+    pos_weight = torch.where(torch.greater_equal(labels, pos_threshold), ones - probs, zeros)
+    neg_weight = torch.where(torch.less(labels, pos_threshold), probs, zeros)
+    '''
+    用于保证数值稳定性，log(sigmoid(x)) = log(1/(1+e^-x) = -log(1+e^-x) = x-x-log(1+e^-x) = x-log(e^x +1)
+    pos_loss = tf.where(tf.less(logits,0),logits-tf.log(tf.exp(logits)+1),tf.log(probs))
+    '''
+    pure_pos_loss = -torch.minimum(logits,logits.new_tensor(0,dtype=logits.dtype))+torch.log(1+torch.exp(-torch.abs(logits)))
+    pos_loss = pure_pos_loss*torch.pow(pos_weight, alpha)
+    if sum:
+        pos_loss = torch.sum(pos_loss)
+    '''
+    用于保证数值稳定性
+    '''
+    pure_neg_loss = F.relu(logits)+torch.log(1+torch.exp(-torch.abs(logits)))
+    neg_loss = torch.pow((1 - labels), beta) * torch.pow(neg_weight, alpha) * pure_neg_loss
+    if sum:
+        neg_loss = torch.sum(neg_loss)
+    loss = (pos_loss + neg_loss) / (num_pos + 1e-4)
+    return loss
+
+def focal_mse_loss_for_head_map(gt_value,pred_value,alpha=0.25,gamma=2):
+    '''
+    alpha: pos loss x alpha, neg loss x (1-alpha)
+    gamma: loss x loss^gamma
+    gt_value: value range [0,1]
+    注：
+    
+    '''
+    assert gt_value.numel()==pred_value.numel(), f"ERROR: unmatch gt_value's shape {gt_value.shape} and pred_value's shape {pred_value.shape}"
+    gt_value = gt_value.float()
+    pred_value = pred_value.float()
+    loss = torch.square(gt_value-pred_value)
+    bin_gt_value = torch.greater(gt_value,0).to(pred_value.dtype)
+    total_nr = bin_gt_value.new_tensor(torch.numel(bin_gt_value))
+    #neg_scale = torch.sum(bin_gt_value)/total_nr
+    #pos_scale = bin_gt_value.new_tensor(1.0)-neg_scale
+    num_pos = torch.sum(bin_gt_value)
+    num_neg = total_nr-num_pos
+
+    neg_scale = total_nr/(num_neg*2.0+2e-1)
+    pos_scale = total_nr/(num_pos*2.0+2e-1)
+
+    auto_t = bin_gt_value*pos_scale+(1-bin_gt_value)*neg_scale
+    loss = loss*auto_t
+
+
+    if alpha is not None and alpha>0:
+        alpha_t = bin_gt_value*alpha+(1-bin_gt_value)*(1-alpha)
+        alpha_t = alpha_t/torch.clip(torch.mean(alpha_t).detach(),min=1e-6)
+        loss = loss*alpha_t
+    if gamma is not None and gamma>0:
+        gamma_t = torch.pow(torch.abs(gt_value-pred_value),gamma)
+        gamma_t = gamma_t/torch.clip(torch.mean(gamma_t).detach(),min=1e-6)
+        loss = loss*gamma_t
+    
+    return loss