对行人reid模型进行压缩的常用方法和loss总结

粘粘粘豆包

已于 2023-07-27 19:16:39 修改

阅读量313

点赞数

文章标签：剪枝 pytorch 算法计算机视觉

于 2023-07-27 19:14:52 首次发布

本文链接：https://blog.csdn.net/qq_38853994/article/details/131781473

版权

对行人reid模型进行压缩的常用方法和loss总结

应用背景
方法以及论文

应用背景

在现实应用场景中，行人reid算法通常跟在检测之后。且当行人较多时，如一个场景有20个人，则当检测处理一张图片时，行人reid需要处理20个目标图像。根据目前nvidia提供的基于jetson NX模块的 benchmark，常用的检测模型yolov3-tiny-416，速度达到546.69FPS。而大多数reid模型基于resnet50, ResNet50_224x224 速度达到887.6FPS, 与检测相比，当处理的画面行人较多时，reid会形成瓶颈。因此，对于reid模型，需要更轻量级的backbone或压缩已有的训练好的模型，提升处理速度。
本文总结了目前本人正在尝试的reid模型压缩方法的原理。主要包括常见的KD方法（overhaul) 和模型剪枝方法（network slimming）的原理和常用loss函数原理。

方法以及论文

reid model compression

论文链接
对比了常用的model prune方法，包括unstructured pruning, L1/L2 norm, network slimming。最终选用unstructured pruning 方法生成剪枝后的模型。在fintune阶段，将原模型作为教师模型，剪枝后的模型作为学生模型，进行KD训练。 reid_model_compress_fig
涉及的loss主要为：KL 散度损失函数，metric learning-based KD loss。

KL divergence loss

KL 散度，衡量两个概率分布(P,Q) 的距离。 $P_s$ 学生模型概率分布输出， $P_t$ 教师模型概率分布输出。
$D_{KL} (P||Q) =- \sum_{x\in\mathbb X}P(x)log(\frac{Q(x)}{P(x)}) \\ L_{KD} = T^2\times D_{KL}(P_s/T, P_t/T)$

# from fastreid/modeling/meta_arch/distiller.py
@staticmethod
def _kldiv(y_s, y_t, t):
    p_s = F.log_softmax(y_s / t, dim=1)
    p_t = F.softmax(y_t / t, dim=1)
    loss = F.kl_div(p_s, p_t, reduction="sum") * (t ** 2) / y_s.shape[0]
    return loss

def jsdiv_loss(self, y_s, y_t, t=16):
    loss = (self._kldiv(y_s, y_t, t) + self._kldiv(y_t, y_s, t)) / 2
    return loss
# 最终返回两个kldiv的平均值jsdiv_loss

metric learning-based KD loss

对于一个batch中的N张图像， $x_i$ 为 $image_i$ ， $F(x_i)$ 为没有归一化的图片特征。
$d_{i,j} = || F(x_i)-F(x_j)||_2\\L_{MetricKD} = \frac{1}{N^2}\sum^N_i\sum^N_j|| d^s_{i,j}-d^t_{i,j}||_2$

def euclidean_dist(x, y):
    m, n = x.size(0), y.size(0)
    xx = torch.pow(x, 2).sum(1, keepdim=True).expand(m, n)
    yy = torch.pow(y, 2).sum(1, keepdim=True).expand(n, m).t()
    dist = xx + yy - 2 * torch.matmul(x, y.t())
    dist = dist.clamp(min=1e-12).sqrt()  # for numerical stability
    return dist
def metricKD_loss(s_embedding, t_embedding ):
  '''
  s_embedding:  (N, dim)
  t_embedding:  (N, dim)
  '''
  s_dist_mat = euclidean_dist(s_embedding, s_embedding)
  t_dist_mat = euclidean_dist(t_embedding, t_embedding)
  # N = s_dist_mat.shape[0]
  loss = F.mse_loss(s_dist_mat, t_dist_mat) 
  ## mse_loss reduction 默认为none时。完成sum和loss/N**2
  return loss

network slimming

论文链接
基于bn层，进行sparsity regularization 训练，使得bn权重具有稀疏性。并根据剪枝比例设置阈值，将部分bn weight 置0，从而改变部分conv层的input/output数量，达到剪枝目的。最后，对于剪枝后的模型进行finetune。
netwok_slimming_fig
本人目前尝试用network slimming 做模型剪枝的原因在于，常用环境中的tensorrt7目前不支持对于稀疏矩阵的优化加速。跟unstructured pruning 方法相比，network slimming 方法显著减少了conv层通道数。对于剪枝后的模型，可以在tensorrt中可以定义更小的conv层，从而可以达到在tensorrt7环境中的加速。

本方法涉及到的loss主要为：sparsity regularization 中的L1 norm。

sparsity regularization

bn层定义为 $\hat{z}=\frac{z_{i n}-\mu_{B}}{\sqrt{\sigma_{B}^{2}+\epsilon}};\;\;\bar{z}_{o u t}=\gamma\hat{z}+\beta$ , bn 层通常跟在conv 之后。可以利用 $\gamma$ 权重进行conv层通道的剪枝。sparsity regularization(稀疏约束)加入对于 $\gamma$ 值的约束, $g(\cdot)$ 为L1 norm。loss 计算方法如下：
$L=\sum_{(x,y)}^{\mathbf{\alpha}}l(f(x,W),y)+\lambda\sum_{\gamma\in\Gamma}^{\mathbf{\alpha}}g(\gamma\mathbf{)}$
运用L1 norm原因：L0 norm表示矩阵中非零元素个数，但是非凸的(non-convex)，不利于训练。L1 norm表示矩阵权重的绝对值，为凸函数且除零点外斜率为常数，可以将一些权重限制为0。L2 norm的斜率在接近0的时候极小，所以只能将权重限制为很小的数，但是无法为0。参考链接

# additional subgradient descent on the sparsity-induced penalty term
def updateBN(self):
    s = 0.0001 #\lambda=0.0001
    for name, m in self.named_modules():
        # if isinstance(m, nn.BatchNorm2d):
        #     if name.endswith("BN"): 
        #         continue ### do not regularize IBN module
        if isinstance(m, nn.BatchNorm2d) and (name.endswith("bn1") or name.endswith("bn2")) and name!='bn1': 
            ## do not add sparsity regu on downsample path and IBN module
            m.weight.grad.data.add_(s*torch.sign(m.weight.data))  # L1
## trianer sample
"""
        self.optimizer.zero_grad()
        losses.backward()  ## compute gradient based on losses
        self.model.backbone.updateBN() ### add bn sparsity regularization gradient
        self.optimizer.step() ## update weights
"""

distill overhaul

论文链接
从四个方面探索模型distillation的问题。分为：1.teacher transform (教师模型的特征转换） 2.student transform (学生模型的特征转换）3.distillation feature position (蒸馏特征位置的选取) 4.distance function (特征距离的计算函数）
在这里插入图片描述
分别提出以下方法：

teacher transform

传统方法会reduce the dimension of the feature vector via teacher transform，会损失教师特征的精度。文章提出，不改变教师特征维度，采用margin ReLU 方法进行过滤，去除特征矩阵中为负值的有害信息，保留为正值的有用信息。此处对应后面distillation feature position 中 pre-ReLU 的设计。
阈值： $m_c = E[F_t^i | F_t^i < 0, i \in C ]$

from scipy.stats import norm
def get_margin_from_BN(bn):
    margin = []
    std = bn.weight.data  ### gamma, std after BN
    mean = bn.bias.data   ### beta, std after BN 
    for (s, m) in zip(std, mean): ### every channel
        s = abs(s.item())
        m = m.item()
        if norm.cdf(-m / s) > 0.001:   
       	"""
       	cdf cumulative distribution function (norm normal distribution)
       	累计分布函数，代表了概率
       	"""
            margin.append(- s * math.exp(- (m / s) ** 2 / 2) / \
                          math.sqrt(2 * math.pi) / norm.cdf(-m / s) + m)
        else:
            margin.append(-3 * s)
    return torch.tensor(margin, dtype=torch.float32, device=mean.device)

student transform

一些传统方法，采用 $T_s = T_t$ ，对学生特征也造成一样的损失。文章提出， $T_s$ 采取 $1\times1$ conv 开扩大学生特征的维度，使其和教师特征维度相同。

def build_feature_connector(t_channel, s_channel):
    C = [nn.Conv2d(s_channel, t_channel, kernel_size=1, stride=1, padding=0, bias=False),
         nn.BatchNorm2d(t_channel)] ### bn illustrated in section 3.3
    ### model weigt init, backpropagation works 
    for m in C:
        if isinstance(m, nn.Conv2d):
            n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
            m.weight.data.normal_(0, math.sqrt(2. / n))
        elif isinstance(m, nn.BatchNorm2d):
            m.weight.data.fill_(1)
            m.bias.data.zero_()
    return nn.Sequential(*C)
"""train sample
1. add T_s
 setattr(self, "connectors_{}".format(i), nn.ModuleList(
        [build_feature_connector(t, s) for t, s in zip(t_channels, s_channels)]))
2. get feature after T_s
 s_feats_connect = getattr(self, "connectors_{}".format(i))[j](s_feats[j])
"""

distill feature position

一些有相同空间大小的层为一个layer group (层组)。提出pre_ReLU, 即在每个layer grop之后，ReLU 之前作为进行对比的特征的位置。

distance function

由于教师模型特征选择的位置是pre_ReLU, 教师特征中包含会被后续ReLU模块滤掉的负值(adverse information)。因此，不用传递所有信息。文章提出 partial L2 distance, 来跳过对于负值信息的蒸馏。

def distillation_loss(source, target, margin):
    target = torch.max(target, margin)
    loss = F.mse_loss(source, target, reduction="none")
    loss = loss * ((source > target) | (target > 0)).float() ### "|"means or
    ##针对target >0;和当 target<=0 时，source>target 的区域计算 L2 loss
    return loss.sum()
"""
source: student feature after T_s
target: teacher feature after T_t
margin: margin ReLU threshold
"""

粘粘粘豆包

关注

0
点赞
踩
1

收藏

觉得还不错? 一键收藏
1
评论
对行人reid模型进行压缩的常用方法和loss总结

本文总结了目前本人正在尝试的reid模型压缩方法的原理。主要包括常见的KD方法（overhaul) 和模型剪枝方法（network slimming）的原理和常用loss函数原理。
复制链接

扫一扫