​

import torch
import torch.nn as nn
import torch.nn.functional as Fdef is_parallel(model):"""Returns True if model is of type DP or DDP."""return isinstance(model, (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel))def de_parallel(model):"""De-parallelize a model: returns single-GPU model if model is of type DP or DDP."""return model.module if is_parallel(model) else modelclass MimicLoss(nn.Module):def __init__(self, channels_s, channels_t):super(MimicLoss, self).__init__()device = 'cuda' if torch.cuda.is_available() else 'cpu'self.mse = nn.MSELoss()def forward(self, y_s, y_t):"""Forward computation.Args:y_s (list): The student model prediction withshape (N, C, H, W) in list.y_t (list): The teacher model prediction withshape (N, C, H, W) in list.Return:torch.Tensor: The calculated loss value of all stages."""assert len(y_s) == len(y_t)losses = []for idx, (s, t) in enumerate(zip(y_s, y_t)):assert s.shape == t.shapelosses.append(self.mse(s, t))loss = sum(losses)return lossclass CWDLoss(nn.Module):"""PyTorch version of `Channel-wise Distillation for Semantic Segmentation.<https://arxiv.org/abs/2011.13256>`_."""def __init__(self, channels_s, channels_t, tau=1.0):super(CWDLoss, self).__init__()self.tau = taudef forward(self, y_s, y_t):"""Forward computation.Args:y_s (list): The student model prediction withshape (N, C, H, W) in list.y_t (list): The teacher model prediction withshape (N, C, H, W) in list.Return:torch.Tensor: The calculated loss value of all stages."""assert len(y_s) == len(y_t)losses = []for idx, (s, t) in enumerate(zip(y_s, y_t)):assert s.shape == t.shapeN, C, H, W = s.shape# normalize in channel diemensionsoftmax_pred_T = F.softmax(t.view(-1, W * H) / self.tau, dim=1)  # [N*C, H*W]logsoftmax = torch.nn.LogSoftmax(dim=1)cost = torch.sum(softmax_pred_T * logsoftmax(t.view(-1, W * H) / self.tau) -softmax_pred_T * logsoftmax(s.view(-1, W * H) / self.tau)) * (self.tau ** 2)losses.append(cost / (C * N))loss = sum(losses)return lossclass MGDLoss(nn.Module):def __init__(self, channels_s, channels_t, alpha_mgd=0.00002, lambda_mgd=0.65):super(MGDLoss, self).__init__()device = 'cuda' if torch.cuda.is_available() else 'cpu'self.alpha_mgd = alpha_mgdself.lambda_mgd = lambda_mgdself.generation = [nn.Sequential(nn.Conv2d(channel, channel, kernel_size=3, padding=1),nn.ReLU(inplace=True),nn.Conv2d(channel, channel, kernel_size=3, padding=1)).to(device) for channel in channels_t]def forward(self, y_s, y_t):"""Forward computation.Args:y_s (list): The student model prediction withshape (N, C, H, W) in list.y_t (list): The teacher model prediction withshape (N, C, H, W) in list.Return:torch.Tensor: The calculated loss value of all stages."""assert len(y_s) == len(y_t)losses = []for idx, (s, t) in enumerate(zip(y_s, y_t)):assert s.shape == t.shapelosses.append(self.get_dis_loss(s, t, idx) * self.alpha_mgd)loss = sum(losses)return lossdef get_dis_loss(self, preds_S, preds_T, idx):loss_mse = nn.MSELoss(reduction='sum')N, C, H, W = preds_T.shapedevice = preds_S.devicemat = torch.rand((N, 1, H, W)).to(device)mat = torch.where(mat > 1 - self.lambda_mgd, 0, 1).to(device)masked_fea = torch.mul(preds_S, mat)new_fea = self.generation[idx](masked_fea)dis_loss = loss_mse(new_fea, preds_T) / Nreturn dis_lossclass Distill_LogitLoss:def __init__(self, p, t_p, alpha=0.25):t_ft = torch.cuda.FloatTensor if t_p[0].is_cuda else torch.Tensorself.p = pself.t_p = t_pself.logit_loss = t_ft([0])self.DLogitLoss = nn.MSELoss(reduction="none")self.bs = p[0].shape[0]self.alpha = alphadef __call__(self):# per outputassert len(self.p) == len(self.t_p)for i, (pi, t_pi) in enumerate(zip(self.p, self.t_p)):  # layer index, layer predictionsassert pi.shape == t_pi.shapeself.logit_loss += torch.mean(self.DLogitLoss(pi, t_pi))return self.logit_loss[0] * self.alphadef get_fpn_features(x, model, fpn_layers=[15, 18, 21]):y, fpn_feats = [], []with torch.no_grad():model = de_parallel(model)module_list = model.model[:-1] if hasattr(model, "model") else model[:-1]for m in module_list:# if not from previous layerif m.f != -1:x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]  # from earlier layersx = m(x)y.append(x if m.i in model.save else None)  # save outputif m.i in fpn_layers:fpn_feats.append(x)return fpn_featsdef get_channels(model, fpn_layers=[15, 18, 21]):y, out_channels = [], []p = next(model.parameters())x = torch.zeros((1, 3, 64, 64), device=p.device)with torch.no_grad():model = de_parallel(model)module_list = model.model[:-1] if hasattr(model, "model") else model[:-1]for m in module_list:# if not from previous layerif m.f != -1:x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]  # from earlier layersx = m(x)y.append(x if m.i in model.save else None)  # save outputif m.i in fpn_layers:out_channels.append(x.shape[1])return out_channelsclass FeatureLoss(nn.Module):def __init__(self, channels_s, channels_t, distiller='cwd'):super(FeatureLoss, self).__init__()device = 'cuda' if torch.cuda.is_available() else 'cpu'self.align_module = nn.ModuleList([nn.Conv2d(channel, tea_channel, kernel_size=1, stride=1, padding=0).to(device)for channel, tea_channel in zip(channels_s, channels_t)])self.norm = [nn.BatchNorm2d(tea_channel, affine=False).to(device)for tea_channel in channels_t]if distiller == 'mimic':self.feature_loss = MimicLoss(channels_s, channels_t)elif distiller == 'mgd':self.feature_loss = MGDLoss(channels_s, channels_t)elif distiller == 'cwd':self.feature_loss = CWDLoss(channels_s, channels_t)else:raise NotImplementedErrordef forward(self, y_s, y_t):assert len(y_s) == len(y_t)tea_feats = []stu_feats = []for idx, (s, t) in enumerate(zip(y_s, y_t)):s = self.align_module[idx](s)s = self.norm[idx](s)t = self.norm[idx](t)tea_feats.append(t)stu_feats.append(s)loss = self.feature_loss(stu_feats, tea_feats)return loss

        #------------------------------Add-Param-Start---------------self.featureloss = 0self.logitloss = 0self.teacherloss = 0self.distillloss =Noneself.model_t = overrides.get("model_t", None)self.distill_feat_type = "cwd"  # "cwd","mgd","mimic"self.distillonline = False  # False or Trueself.logit_loss = False # False or Trueself.distill_layers = [2, 4, 6, 8, 12, 15, 18, 21]# ------------------------------Add-Param-End-----------------

        if self.model_t is not None:for k, v in self.model_t.model.named_parameters():v.requires_grad = Trueself.model_t = self.model_t.to(self.device)

            if self.model_t is not None:self.model_t = nn.parallel.DistributedDataParallel(self.model_t, device_ids=[RANK])

        self.optimizer = self.build_optimizer(model=self.model,model_t=self.model_t,distillloss=self.distillloss,distillonline=self.distillonline,name=self.args.optimizer,lr=self.args.lr0,momentum=self.args.momentum,decay=weight_decay,iterations=iterations)

        self.model = de_parallel(self.model)if self.model_t is not None:self.model_t = de_parallel(self.model_t)self.channels_s = get_channels(self.model,self.distill_layers)self.channels_t = get_channels(self.model_t,self.distill_layers)self.distillloss = FeatureLoss(channels_s=self.channels_s, channels_t=self.channels_t, distiller= self.distill_feat_type)

            if self.model_t is not None:self.model_t.eval()

                    if self.model_t is not None:distill_weight = ((1 - math.cos(i * math.pi / len(self.train_loader))) / 2) * (0.1 - 1) + 1with torch.no_grad():pred_t_offline = self.model_t(batch['img'])tea_features = get_fpn_features(batch['img'], self.model_t,fpn_layers=self.distill_layers)  # forwardself.featureloss = self.distillloss(stu_features, tea_features) * distill_weightself.loss += self.featurelossif self.distillonline:self.model_t.train()pred_t_online = self.model_t(batch['img'])for p in pred_t_online:p = p.detach()if i == 0 and epoch == 0:self.model_t.args["box"] = self.model.args.boxself.model_t.args["cls"] = self.model.args.clsself.model_t.args["dfl"] = self.model.args.dflself.model_t.args = IterableSimpleNamespace(**self.model_t.args)self.teacherloss, _ = self.model_t(batch, pred_t_online)if RANK != -1:self.teacherloss *= world_sizeself.loss += self.teacherlossif self.logit_loss:if not self.distillonline:distill_logit = Distill_LogitLoss(pred_s, pred_t_offline)else:distill_logit = Distill_LogitLoss(pred_s, pred_t_online)self.logitloss = distill_logit()self.loss += self.logitloss

                    pbar.set_description(('%12s' * 2 + '%12.4g' * (5 + loss_len)) %(f'{epoch + 1}/{self.epochs}', mem, *losses, self.featureloss, self.teacherloss, self.logitloss, batch['cls'].shape[0], batch['img'].shape[-1]))

, model_t, distillloss, distillonline=False,

        if model_t is not None and distillonline:for v in model_t.modules():# print(v)if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):  # bias (no decay)g[2].append(v.bias)if isinstance(v, bn):  # weight (no decay)g[1].append(v.weight)elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):  # weight (with decay)g[0].append(v.weight)if model_t is not None and distillloss is not None:for k, v in distillloss.named_modules():# print(v)if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):  # bias (no decay)g[2].append(v.bias)if isinstance(v, bn) or 'bn' in k:  # weight (no decay)g[1].append(v.weight)elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):  # weight (with decay)g[0].append(v.weight)

        return ('\n' + '%12s' *(7 + len(self.loss_names))) % ('Epoch', 'GPU_mem', *self.loss_names, 'dfeaLoss', 'dlineLoss', 'dlogitLoss', 'Instances','Size')

import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLOif __name__ == '__main__':model_t = YOLO(r'weights/yolov8l.pt')  # 此处填写教师模型的权重文件地址model_t.model.model[-1].set_Distillation = True  # 不用理会此处用于设置模型蒸馏model_s = YOLO(r'ultralytics/cfg/models/v8/yolov8.yaml')  # 学生文件的yaml文件 or 权重文件地址model_s.train(data=r'C:\Users\Administrator\Desktop\Snu77\ultralytics-main\New_GC-DET\data.yaml',  #  将data后面替换你自己的数据集地址cache=False,imgsz=640,epochs=100,single_cls=False,  # 是否是单类别检测batch=1,close_mosaic=10,workers=0,device='0',optimizer='SGD',  # using SGDamp=True,  # 如果出现训练损失为Nan可以关闭ampproject='runs/train',name='exp',model_t=model_t.model)