yolov7改进优化之蒸馏（二）

yolov7改进优化之蒸馏（二）

续yolov7改进优化之蒸馏（一）-CSDN博客
上一篇已经基本写出来yolov7/v5蒸馏的整个过程，不过要真的训起来我们还需要进行一些修改。

Model修改

蒸馏需要对teacher和student网络的特征层进行loss计算，因此我们forward时要能够返回需要的中间层，这需要修改yolo.py中的Model类。

forward_once接口修改

增加接口参数 extra_features用于指定要返回的中间层的索引:

def forward_once(self, x, profile=False, extra_features: list = []):y, dt = [], []  # outputsfeatures = []for i, m in enumerate(self.model):if m.f != -1:  # if not from previous layerx = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]  # from earlier layersif not hasattr(self, "traced"):self.traced = Falseif self.traced:if (isinstance(m, Detect)or isinstance(m, IDetect)or isinstance(m, IAuxDetect)or isinstance(m, IKeypoint)):breakif profile:c = isinstance(m, (Detect, IDetect, IAuxDetect, IBin))o = thop.profile(m, inputs=(x.copy() if c else x,), verbose=False)[0] / 1e9 * 2 if thop else 0  # FLOPSfor _ in range(10):m(x.copy() if c else x)t = time_synchronized()for _ in range(10):m(x.copy() if c else x)dt.append((time_synchronized() - t) * 100)print("%10.1f%10.0f%10.1fms %-40s" % (o, m.np, dt[-1], m.type))x = m(x)  # runy.append(x if m.i in self.save else None)  # save outputif i in extra_features:features.append(x)if not self.training and len(extra_features) != 0 and len(extra_features) == len(features):return x, featuresif profile:print("%.1fms total" % sum(dt))if len(extra_features) != 0:return x, featuresif self.training and isinstance(x, tuple):x = x[-1]return x

主要增加将中间层返回的代码。

forward接口修改

forward接口调用了forward_once接口，因此，forward接口也需要增加这个参数。

def forward(self, x, augment=False, profile=False, extra_features: list = []):if augment:img_size = x.shape[-2:]  # height, widths = [1, 0.83, 0.67]  # scalesf = [None, 3, None]  # flips (2-ud, 3-lr)y = []  # outputsfor si, fi in zip(s, f):xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))yi = self.forward_once(xi)[0]  # forward# cv2.imwrite(f'img_{si}.jpg', 255 * xi[0].cpu().numpy().transpose((1, 2, 0))[:, :, ::-1])  # saveyi[..., :4] /= si  # de-scaleif fi == 2:yi[..., 1] = img_size[0] - yi[..., 1]  # de-flip udelif fi == 3:yi[..., 0] = img_size[1] - yi[..., 0]  # de-flip lry.append(yi)return torch.cat(y, 1), None  # augmented inference, trainelse:return self.forward_once(x, profile, extra_features)  # single-scale inference, train

hyp文件修改

在hyp文件中添加student_kd_layers和teacher_kd_layers来指定要蒸馏的层，我们可以指定IDetect前面的三个特征层：

student_kd_layers: [75,88,101]
teacher_kd_layers: [75,88,101]

训练

训练方式与正常训练一样，只是启动时要指定teacher-weights。

结语

这一篇结合上一篇就可以吧基于FGD算法的蒸馏训练起来了，其他蒸馏的修改也大同小异了。