人物属性模型移动端实验记录

class KLSoftLoss(nn.Module):
    r"""Apply softtarget for kl loss

    Arguments:
        reduction (str): "batchmean" for the mean loss with the p(x)*(log(p(x)) - log(q(x)))
    """
    def __init__(self, temperature=1, reduction="batchmean"):
        super(KLSoftLoss, self).__init__()
        self.reduction = reduction
        self.eps = 1e-7
        self.temperature = temperature
        self.klloss = nn.KLDivLoss(reduction=self.reduction)

    def forward(self, s_logits, t_logits):
        s_prob = F.log_softmax(s_logits / self.temperature, 1)
        t_prob = F.softmax(t_logits / self.temperature, 1)
        loss = self.klloss(s_prob, t_prob) * self.temperature * self.temperature
        return loss

  optimizer.zero_grad()
  loss.backward()

  # use the slimming prune for training
  if args.prune and args.use_sr:
      for m in model.modules():
          if isinstance(m, nn.BatchNorm2d):
              m.weight.grad.data.add_(args.sr * torch.sign(m.weight.data))

  optimizer.step()

def prune_only_res_hidden(percent, model, keep_channel=True, channel_ratio=8, cuda=True):
    """only prune the inverResidual module first bn layer
    """
    total = 0
    highest_thre = []
    for m in model.modules():
        if isinstance(m, InvertedResidual):
            # only prune the 3 conv layer
            if len(m.conv) > 5:
                for i in range(len(m.conv)):
                    if i == 1:
                        if isinstance(m.conv[i], nn.BatchNorm2d):
                            total += m.conv[i].weight.data.shape[0]
                            highest_thre.append(m.conv[i].weight.data.abs().max().item())
                            total += m.conv[i+3].weight.data.shape[0]
                            highest_thre.append(m.conv[i+3].weight.data.abs().max().item())



    bn = torch.zeros(total)
    index = 0
    for m in model.modules():
        if isinstance(m, InvertedResidual):
            # only prune the 3 conv layer
            if len(m.conv) > 5:
                for i in range(len(m.conv)):
                    if i != len(m.conv) - 1:
                        if isinstance(m.conv[i], nn.BatchNorm2d):
                            size = m.conv[i].weight.data.shape[0]
                            bn[index:(index+size)] = m.conv[i].weight.data.abs().clone()
                            index += size

    print(bn.size())
    y, i = torch.sort(bn)
    thre_index = int(total * percent)
    thre = y[thre_index]
    highest_thre = min(highest_thre)

    # 判断阈值
    if thre > highest_thre:
        thre = highest_thre

    print("the min thre is {}, the max thre is {}!!!!".format(thre, highest_thre))
    pruned = 0
    c = {}
    cfg_mask = []
    idx = 0

    for m in model.modules():
        if isinstance(m, InvertedResidual):
            # only prune the 3 conv layer
            if len(m.conv) > 5:
                for i in range(len(m.conv)):
                    if i == 1:
                        if isinstance(m.conv[i], nn.BatchNorm2d):
                            weight_copy = m.conv[i].weight.data.clone()
                            if cuda:
                                mask = weight_copy.abs().gt(thre).float().cuda()
                            else:
                                mask = weight_copy.abs().gt(thre).float()

                            if keep_channel:
                                keep_channel_number = get_min_number(torch.sum(mask), channel_ratio)
                                if torch.sum(mask) < keep_channel_number:
                                    n = int(keep_channel_number - torch.sum(mask))
                                    mask_index = torch.where(mask==0)[0]
                                    new_weight = weight_copy.abs()[mask_index]
                                    _, weight_index = torch.sort(new_weight)
                                    w_index = mask_index[weight_index[-n: ]]
                                    mask[w_index] = 1.0

                            pruned = pruned + mask.shape[0] - torch.sum(mask)
                            # first conv + bn
                            m.conv[i].weight.data.mul_(mask)
                            m.conv[i].bias.data.mul_(mask)
                            # second conv + bn
                            m.conv[i+3].weight.data.mul_(mask)
                            m.conv[i+3].bias.data.mul_(mask)
                            c[idx] = int(torch.sum(mask))
                            cfg_mask.append(mask.clone())

                            print('layer index: {:d} \t total channel: {:d} \t remaining channel: {:d}'.format(idx, mask.shape[0], int(torch.sum(mask))))
                            idx += 1
    print(c)
    print(len(c))
    print(len(cfg_mask))
    # pruned_ratio = pruned / total
    print('Pre-processing Successful!!!')
    return model, cfg_mask, c