yolov7改进优化之蒸馏（一）

最近比较忙，有一段时间没更新了，最近yolov7用的比较多，总结一下。上一篇yolov5及yolov7实战之剪枝_CodingInCV的博客-CSDN博客 我们讲了通过剪枝来裁剪我们的模型，达到在精度损失不大的情况下，提高模型速度的目的。上一篇是从速度的角度，这一篇我们从检测性能的角度来改进yolov7（yolov5也类似）。

对于提高检测器的性能，我们除了可以从增加数据、修改模型结构、修改loss等模型本身的角度出发外，深度学习领域还有一个方式---蒸馏。简单的说，蒸馏就是让性能更强的模型（teacher, 参数量更大）来指导性能更弱student模型，从而提高student模型的性能。

蒸馏的方式有很多种，比较简单暴力的比如直接让student模型来拟合teacher模型的输出特征图，当然蒸馏也不是万能的，毕竟student模型和teacher模型的参数量有差距，student模型不一定能很好的学习teacher的知识，对于自己的任务有没有作用也需要尝试。

本篇选择的方法是去年CVPR上的针对目标检测的蒸馏算法：
yzd-v/FGD: Focal and Global Knowledge Distillation for Detectors (CVPR 2022) (github.com)

针对该方法的解读可以参考：FGD-CVPR2022：针对目标检测的焦点和全局蒸馏 - 知乎 (zhihu.com)

本篇暂时不涉及理论，重点在把这个方法集成到yolov7训练。步骤如下。

载入teacher模型

蒸馏首先需要有一个teacher模型，这个teacher模型一般和student同样结构，只是参数量更大、层数更多。比如对于yolov5，可以尝试用yolov5m来蒸馏yolov5s。

train.py增加一个命令行参数：

    parser.add_argument("--teacher-weights", type=str, default="", help="initial weights path")

在train函数中载入teacher weights，过程与原有的载入过程类似，注意，DP或者DDP模型也要对teacher模型做对应的处理。

# teacher model
    if opt.teacher_weights:
        teacher_weights = opt.teacher_weights
        # with torch_distributed_zero_first(rank):
        #     teacher_weights = attempt_download(teacher_weights)  # download if not found locally
        teacher_model = Model(teacher_weights, ch=3, nc=nc).to(device)  # create    
        # load state_dict
        ckpt = torch.load(teacher_weights, map_location=device)  # load checkpoint
        state_dict = ckpt["model"].float().state_dict()  # to FP32
        teacher_model.load_state_dict(state_dict, strict=True)  # load
        #set to eval
        teacher_model.eval()
        #set IDetect to train mode
        # teacher_model.model[-1].train()
        logger.info(f"Load teacher model from {teacher_weights}")  # report

    # DP mode
    if cuda and rank == -1 and torch.cuda.device_count() > 1:
        model = torch.nn.DataParallel(model)
        if opt.teacher_weights:
            teacher_model = torch.nn.DataParallel(teacher_model)
            
	 # SyncBatchNorm
    if opt.sync_bn and cuda and rank != -1:
        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
        logger.info("Using SyncBatchNorm()")
        if opt.teacher_weights:
	        teacher_model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(teacher_model).to(device)

teacher模型不进行梯度计算，因此：

if opt.teacher_weights:
        for param in teacher_model.parameters():
            param.requires_grad = False

蒸馏Loss

蒸馏loss是计算teacher模型的一层或者多层与student的对应层的相似度，监督student模型向teacher模型靠近。对于yolov7，可以去监督三个特征层。

参考FGD的开源代码，我们在loss.py中增加一个FeatureLoss类, 参数暂时使用默认：

class FeatureLoss(nn.Module):

    """PyTorch version of `Feature Distillation for General Detectors`
   
    Args:
        student_channels(int): Number of channels in the student's feature map.
        teacher_channels(int): Number of channels in the teacher's feature map. 
        temp (float, optional): Temperature coefficient. Defaults to 0.5.
        name (str): the loss name of the layer
        alpha_fgd (float, optional): Weight of fg_loss. Defaults to 0.001
        beta_fgd (float, optional): Weight of bg_loss. Defaults to 0.0005
        gamma_fgd (float, optional): Weight of mask_loss. Defaults to 0.0005
        lambda_fgd (float, optional): Weight of relation_loss. Defaults to 0.000005
    """
    def __init__(self,
                 student_channels,
                 teacher_channels,
                 temp=0.5,
                 alpha_fgd=0.001,
                 beta_fgd=0.0005,
                 gamma_fgd=0.001,
                 lambda_fgd=0.000005,
                 ):
        super(FeatureLoss, self).__init__()
        self.temp = temp
        self.alpha_fgd = alpha_fgd
        self.beta_fgd = beta_fgd
        self.gamma_fgd = gamma_fgd
        self.lambda_fgd = lambda_fgd
    
        if student_channels != teacher_channels:
            self.align = nn.Conv2d(student_channels, teacher_channels, kernel_size=1, stride=1, padding=0)
        else:
            self.align = None
        
        self.conv_mask_s = nn.Conv2d(teacher_channels, 1, kernel_size=1)
        self.conv_mask_t = nn.Conv2d(teacher_channels, 1, kernel_size=1)
        self.channel_add_conv_s = nn.Sequential(
            nn.Conv2d(teacher_channels, teacher_channels//2, kernel_size=1),
            nn.LayerNorm([teacher_channels//2, 1, 1]),
            nn.ReLU(inplace=True),  # yapf: disable
            nn.Conv2d(teacher_channels//2, teacher_channels, kernel_size=1))
        self.channel_add_conv_t = nn.Sequential(
            nn.Conv2d(teacher_channels, teacher_channels//2, kernel_size=1),
            nn.LayerNorm([teacher_channels//2, 1, 1]),
            nn.ReLU(inplace=True),  # yapf: disable
            nn.Conv2d(teacher_channels//2, teacher_channels, kernel_size=1))

        self.reset_parameters()

    def forward(self,
                preds_S,
                preds_T,
                gt_bboxes,
                img_metas):
        """Forward function.
        Args:
            preds_S(Tensor): Bs*C*H*W, student's feature map
            preds_T(Tensor): Bs*C*H*W, teacher's feature map
            gt_bboxes(tuple): Bs*[nt*4], pixel decimal: (tl_x, tl_y, br_x, br_y)
            img_metas (list[dict]): Meta information of each image, e.g.,
            image size, scaling factor, etc.
        """
        assert preds_S.shape[-2:] == preds_T.shape[-2:], 'the output dim of teacher and student differ'
        device = gt_bboxes.device
        self.to(device)
        if self.align is not None:
            preds_S = self.align(preds_S)

        N,C,H,W = preds_S.shape

        S_attention_t, C_attention_t = self.get_attention(preds_T, self.temp)
        S_attention_s, C_attention_s = self.get_attention(preds_S, self.temp)
        
        Mask_fg = torch.zeros_like(S_attention_t)
        # Mask_bg = torch.ones_like(S_attention_t)
        wmin,wmax,hmin,hmax = [],[],[],[]
        img_h, img_w = img_metas
        bboxes = gt_bboxes[:,2:6]
        #xywh2xyxy
        bboxes = xywh2xyxy(bboxes)
        new_boxxes = torch.ones_like(bboxes)
        new_boxxes[:, 0] = torch.floor(bboxes[:, 0]*W)
        new_boxxes[:, 2] = torch.ceil(bboxes[:, 2]*W)
        new_boxxes[:, 1] = torch.floor(bboxes[:, 1]*H)
        new_boxxes[:, 3] = torch.ceil(bboxes[:, 3]*H)

        #to int
        new_boxxes = new_boxxes.int()

        for i in range(N):
            new_boxxes_i = new_boxxes[torch.where(gt_bboxes[:,0]==i)]

            wmin.append(new_boxxes_i[:, 0])
            wmax.append(new_boxxes_i[:, 2])
            hmin.append(new_boxxes_i[:, 1])
            hmax.append(new_boxxes_i[:, 3])

            area = 1.0/(hmax[i].view(1,-1)+1-hmin[i].view(1,-1))/(wmax[i].view(1,-1)+1-wmin[i].view(1,-1))

            for j in range(len(new_boxxes_i)):
                Mask_fg[i][hmin[i][j]:hmax[i][j]+1, wmin[i][j]:wmax[i][j]+1] = \
                        torch.maximum(Mask_fg[i][hmin[i][j]:hmax[i][j]+1, wmin[i][j]:wmax[i][j]+1], area[0][j])

        Mask_bg = torch.where(Mask_fg > 0, 0., 1.)
        Mask_bg_sum = torch.sum(Mask_bg, dim=(1,2))
        Mask_bg[Mask_bg_sum>0] /= Mask_bg_sum[Mask_bg_sum>0].unsqueeze(1).unsqueeze(2)

        fg_loss, bg_loss = self.get_fea_loss(preds_S, preds_T, Mask_fg, Mask_bg, 
                        C_attention_s, C_attention_t, S_attention_s, S_attention_t)
        mask_loss = self.get_mask_loss(C_attention_s, C_attention_t, S_attention_s, S_attention_t)
        rela_loss = self.get_rela_loss(preds_S, preds_T)

        loss = self.alpha_fgd * fg_loss + self.beta_fgd * bg_loss \
            + self.gamma_fgd * mask_loss + self.lambda_fgd * rela_loss
            
        return loss, loss.detach()

    def get_attention(self, preds, temp):
        """ preds: Bs*C*W*H """
        N, C, H, W= preds.shape

        value = torch.abs(preds)
        # Bs*W*H
        fea_map = value.mean(axis=1, keepdim=True)
        S_attention = (H * W * F.softmax((fea_map/temp).view(N,-1), dim=1)).view(N, H, W)

        # Bs*C
        channel_map = value.mean(axis=2,keepdim=False).mean(axis=2,keepdim=False)
        C_attention = C * F.softmax(channel_map/temp, dim=1)

        return S_attention, C_attention


    def get_fea_loss(self, preds_S, preds_T, Mask_fg, Mask_bg, C_s, C_t, S_s, S_t):
        loss_mse = nn.MSELoss(reduction='sum')
        
        Mask_fg = Mask_fg.unsqueeze(dim=1)
        Mask_bg = Mask_bg.unsqueeze(dim=1)

        C_t = C_t.unsqueeze(dim=-1)
        C_t = C_t.unsqueeze(dim=-1)

        S_t = S_t.unsqueeze(dim=1)

        fea_t= torch.mul(preds_T, torch.sqrt(S_t))
        fea_t = torch.mul(fea_t, torch.sqrt(C_t))
        fg_fea_t = torch.mul(fea_t, torch.sqrt(Mask_fg))
        bg_fea_t = torch.mul(fea_t, torch.sqrt(Mask_bg))

        fea_s = torch.mul(preds_S, torch.sqrt(S_t))
        fea_s = torch.mul(fea_s, torch.sqrt(C_t))
        fg_fea_s = torch.mul(fea_s, torch.sqrt(Mask_fg))
        bg_fea_s = torch.mul(fea_s, torch.sqrt(Mask_bg))

        fg_loss = loss_mse(fg_fea_s, fg_fea_t)/len(Mask_fg)
        bg_loss = loss_mse(bg_fea_s, bg_fea_t)/len(Mask_bg)

        return fg_loss, bg_loss


    def get_mask_loss(self, C_s, C_t, S_s, S_t):

        mask_loss = torch.sum(torch.abs((C_s-C_t)))/len(C_s) + torch.sum(torch.abs((S_s-S_t)))/len(S_s)

        return mask_loss
     
    
    def spatial_pool(self, x, in_type):
        batch, channel, width, height = x.size()
        input_x = x
        # [N, C, H * W]
        input_x = input_x.view(batch, channel, height * width)
        # [N, 1, C, H * W]
        input_x = input_x.unsqueeze(1)
        # [N, 1, H, W]
        if in_type == 0:
            context_mask = self.conv_mask_s(x)
        else:
            context_mask = self.conv_mask_t(x)
        # [N, 1, H * W]
        context_mask = context_mask.view(batch, 1, height * width)
        # [N, 1, H * W]
        context_mask = F.softmax(context_mask, dim=2)
        # [N, 1, H * W, 1]
        context_mask = context_mask.unsqueeze(-1)
        # [N, 1, C, 1]
        context = torch.matmul(input_x, context_mask)
        # [N, C, 1, 1]
        context = context.view(batch, channel, 1, 1)

        return context


    def get_rela_loss(self, preds_S, preds_T):
        loss_mse = nn.MSELoss(reduction='sum')

        context_s = self.spatial_pool(preds_S, 0)
        context_t = self.spatial_pool(preds_T, 1)

        out_s = preds_S
        out_t = preds_T

        channel_add_s = self.channel_add_conv_s(context_s)
        out_s = out_s + channel_add_s

        channel_add_t = self.channel_add_conv_t(context_t)
        out_t = out_t + channel_add_t

        rela_loss = loss_mse(out_s, out_t)/len(out_s)
        
        return rela_loss


    def last_zero_init(self, m):
        if isinstance(m, nn.Sequential):
            constant_init(m[-1], val=0)
        else:
            constant_init(m, val=0)

    
    def reset_parameters(self):
        kaiming_init(self.conv_mask_s, mode='fan_in')
        kaiming_init(self.conv_mask_t, mode='fan_in')
        self.conv_mask_s.inited = True
        self.conv_mask_t.inited = True

        self.last_zero_init(self.channel_add_conv_s)
        self.last_zero_init(self.channel_add_conv_t)

实例化FeatureLoss

在train.py中，实例化我们定义的FeatureLoss，由于我们要蒸馏三层，所以需要定一个蒸馏损失的数组：

if opt.teacher_weights:
        student_kd_layers = hyp["student_kd_layers"]
        teacher_kd_layers = hyp["teacher_kd_layers"]
        dump_image = torch.zeros((1, 3, imgsz, imgsz), device=device)
        targets = torch.Tensor([[0, 0, 0, 0, 0, 0]]).to(device)
        _, features = model(dump_image, extra_features = student_kd_layers)  # forward
        _, teacher_features = teacher_model(dump_image,
                                               extra_features=teacher_kd_layers)
        kd_losses = []
        for i in range(len(features)):
            feature = features[i]
            teacher_feature = teacher_features[i]
            _, student_channels, _ , _ = feature.shape
            _, teacher_channels, _ , _ = teacher_feature.shape

            kd_losses.append(FeatureLoss(student_channels,teacher_channels))

其中hyp['xxx_kd_layers']是用于指定我们要蒸馏的层序号。

为了提取出我们需要的层的特征图，我们还需要对模型推理的代码进行修改，这个放在下一篇，这一篇先把主要流程过一遍。

蒸馏训练

与普通loss一样，在训练中，首先计算蒸馏loss, 然后进行反向传播，区别只是计算蒸馏loss时需要使用teacher模型也对数据进行推理。

if opt.teacher_weights:
	pred, features = model(imgs, extra_features = student_kd_layers)  # forward
	_, teacher_features = teacher_model(imgs, extra_features = teacher_kd_layers)
	if "loss_ota" not in hyp or hyp["loss_ota"] == 1 and epoch >= ota_start:
		loss, loss_items = compute_loss_ota(
			pred, targets.to(device), imgs
		)
	else:
		loss, loss_items = compute_loss(
			pred, targets.to(device)
		)  # loss scaled by batch_size
	# kd loss
	loss_items = torch.cat((loss_items[0].unsqueeze(0), loss_items[1].unsqueeze(0), loss_items[2].unsqueeze(0), torch.zeros(1, device=device), loss_items[3].unsqueeze(0)))
	loss_items[-1]*=imgs.shape[0]
	for i in range(len(features)):
		feature = features[i]
		teacher_feature = teacher_features[i]

		kd_loss, kd_loss_item = kd_losses[i](feature, teacher_feature, targets.to(device), [imgsz,imgsz])
		loss += kd_loss
		loss_items[3] += kd_loss_item
		loss_items[4] += kd_loss_item

在这里，我们将kd_loss累加到了loss上。计算出总的loss,其他就与普通训练一样了。

结语

这篇文章简述了一下yolov7的蒸馏过程，更多细节将在下一篇中讲述。