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 = [], [] # outputs
features = []
for i, m in enumerate(self.model):
if m.f != -1: # if not from previous layer
x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers
if not hasattr(self, "traced"):
self.traced = False
if self.traced:
if (
isinstance(m, Detect)
or isinstance(m, IDetect)
or isinstance(m, IAuxDetect)
or isinstance(m, IKeypoint)
):
break
if 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 # FLOPS
for _ 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) # run
y.append(x if m.i in self.save else None) # save output
if i in extra_features:
features.append(x)
if not self.training and len(extra_features) != 0 and len(extra_features) == len(features):
return x, features
if profile:
print("%.1fms total" % sum(dt))
if len(extra_features) != 0:
return x, features
if 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, width
s = [1, 0.83, 0.67] # scales
f = [None, 3, None] # flips (2-ud, 3-lr)
y = [] # outputs
for 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]) # save
yi[..., :4] /= si # de-scale
if fi == 2:
yi[..., 1] = img_size[0] - yi[..., 1] # de-flip ud
elif fi == 3:
yi[..., 0] = img_size[1] - yi[..., 0] # de-flip lr
y.append(yi)
return torch.cat(y, 1), None # augmented inference, train
else:
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算法的蒸馏训练起来了，其他蒸馏的修改也大同小异了。
