Browse Source

train CSPDarkNet-M

yjh0410 2 years ago
parent
commit
60018cbba5
2 changed files with 2 additions and 15 deletions
  1. 1 1
      README.md
  2. 1 14
      models/detectors/__init__.py

+ 1 - 1
README.md

@@ -120,7 +120,7 @@ python train.py --cuda -d coco --root path/to/COCO -v yolov1 -bs 16 --max_epoch
 
 - *All models are trained with ImageNet pretrained weight (IP). All FLOPs are measured with a 640x640 image size on COCO val2017. The FPS is measured with batch size 1 on 3090 GPU from the model inference to the NMS operation.*
 
-- *The reproduced YOLOv5's head is the **Decoupled Head**, which is why the FLOPs and Params are higher than the official YOLOv5.*
+- *The reproduced YOLOv5's head is the **Decoupled Head**, which is why the FLOPs and Params are higher than the official YOLOv5. Due to my limited computing resources, I can not align the training configuration with the official YOLOv5, so I cannot fully replicate the official performance. The YOLOv5 I reproduce is for learning purposes only.*
 
 - *Due to my limited computing resources, I had to abandon training on other YOLO detectors, including YOLOv7-Huge and YOLOv8-Nano~Large. If you are interested in these models and have trained them using the code from this project, I would greatly appreciate it if you could share the trained weight files with me.*
 

+ 1 - 14
models/detectors/__init__.py

@@ -80,21 +80,8 @@ def build_model(args,
             checkpoint = torch.load(args.resume, map_location='cpu')
             # checkpoint state dict
             checkpoint_state_dict = checkpoint.pop("model")
-            # check
-            new_checkpoint_state_dict = {}
+            model.load_state_dict(checkpoint_state_dict)
 
-            for k in list(checkpoint_state_dict.keys()):
-                v = checkpoint_state_dict[k]
-                if 'reduce_layer_3' in k:
-                    k_new = k.split('.')
-                    k_new[1] = 'downsample_layer_1'
-                    k = k_new[0] + '.' + k_new[1] + '.' + k_new[2] + '.' + k_new[3] + '.' + k_new[4]
-                elif 'reduce_layer_4' in k:
-                    k_new = k.split('.')
-                    k_new[1] = 'downsample_layer_2'
-                    k = k_new[0] + '.' + k_new[1] + '.' + k_new[2] + '.' + k_new[3] + '.' + k_new[4]
-                new_checkpoint_state_dict[k] = v
-            model.load_state_dict(new_checkpoint_state_dict)
         return model, criterion
 
     else: