修改对齐卷积加入优化器

train.py

@@ -8,7 +8,7 @@
 # modelL = YOLO('/home/huangzm/code/mycode/pytorch/ultralytics/runs/detect/coco_v8l/weights/best.pt')
 # model_t = YOLO('/home/huangzm/code/mycode/pytorch/ultralytics/runs/detect/train37/weights/best.pt')
 
-# model_t = YOLO('yolov8l.pt')
+model_t = YOLO('yolov8l.pt')
 # success = modelL.export(format="onnx",device="cpu")
 
 data = "/home/huangzm/code/mycode/pytorch/yolov8/Knowledge_Distillation/ultralytics/datasets/coco.yaml"
@@ -22,8 +22,8 @@
 # success = modeln.export(format="onnx")
 # modelL.val(data=data)
 
-# model_s.train(data=data, epochs=100, imgsz=640, Distillation = model_t.model)
-model_s.train(data=data, epochs=100, imgsz=640, Distillation = None)
+model_s.train(data=data, epochs=100, imgsz=640, Distillation = model_t.model)
+# model_s.train(data=data, epochs=100, imgsz=640, Distillation = None)
 
 
 

ultralytics/yolo/engine/trainer.py

@@ -175,6 +175,8 @@ def __init__(self, modeln,modelL,distiller="CWDLoss"):  # model must be de-paral
         # channels_t=[128,256,512,512,512,256,512,512]
         channels_s=[128,256,128,64,128,256]
         channels_t=[512,512,512,256,512,512]
+        # channels_s=[64,128,256]
+        # channels_t=[256,512,512]
         self.D_loss_fn = FeatureLoss(channels_s=channels_s,channels_t=channels_t)
 
         self.teacher_module_pairs = []
@@ -457,6 +459,7 @@ def _setup_train(self, world_size):
         self.accumulate = max(round(self.args.nbs / self.batch_size), 1)  # accumulate loss before optimizing
         weight_decay = self.args.weight_decay * self.batch_size * self.accumulate / self.args.nbs  # scale weight_decay
         self.optimizer = self.build_optimizer(model=self.model,
+                                              model1=self.Distillation,
                                               name=self.args.optimizer,
                                               lr=self.args.lr0,
                                               momentum=self.args.momentum,
@@ -842,7 +845,7 @@ def resume_training(self, ckpt):
                 self.train_loader.dataset.close_mosaic(hyp=self.args)
 
     @staticmethod
-    def build_optimizer(model, name='Adam', lr=0.001, momentum=0.9, decay=1e-5):
+    def build_optimizer(model,model1, name='Adam', lr=0.001, momentum=0.9, decay=1e-5):
         """
         Builds an optimizer with the specified parameters and parameter groups.
 
@@ -866,6 +869,15 @@ def build_optimizer(model, name='Adam', lr=0.001, momentum=0.9, decay=1e-5):
             elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):  # weight (with decay)
                 g[0].append(v.weight)
 
+        for v in model1.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 name == 'Adam':
             optimizer = torch.optim.Adam(g[2], lr=lr, betas=(momentum, 0.999))  # adjust beta1 to momentum
         elif name == 'AdamW':