keep training YOLOv5-L from 211 epoch

README.md

@@ -113,10 +113,10 @@ python train.py --cuda -d coco --root path/to/COCO -v yolov1 -bs 16 --max_epoch
 
 | Model         |   Backbone         | Scale | Epoch |  FPS  | AP<sup>val<br>0.5:0.95 | AP<sup>val<br>0.5 | FLOPs<br><sup>(G) | Params<br><sup>(M) | Weight |
 |---------------|--------------------|-------|-------|-------|------------------------|-------------------|-------------------|--------------------|--------|
-| YOLOv5-N      | CSPDarkNet-N       |  640  |  300  |       |                        |                   |   7.7             |   2.4              |  |
-| YOLOv5-S      | CSPDarkNet-S       |  640  |  300  |       |                        |                   |   27.1            |   9.0              |  |
-| YOLOv5-M      | CSPDarkNet-M       |  640  |  300  |       |                        |                   |   74.3            |   25.4             |  |
-| YOLOv5-L      | CSPDarkNet-L       |  640  |  300  |       |                        |                   |   155.6           |   54.2             |  |
+| YOLOv5-N      | CSPDarkNet-N       |  640  |  250  |       |                        |                   |   7.7             |   2.4              |  |
+| YOLOv5-S      | CSPDarkNet-S       |  640  |  250  |       |                        |                   |   27.1            |   9.0              |  |
+| YOLOv5-M      | CSPDarkNet-M       |  640  |  250  |       |                        |                   |   74.3            |   25.4             |  |
+| YOLOv5-L      | CSPDarkNet-L       |  640  |  250  |       |                        |                   |   155.6           |   54.2             |  |
 
 - *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.*
 

models/detectors/__init__.py

@@ -80,8 +80,21 @@ def build_model(args,
             checkpoint = torch.load(args.resume, map_location='cpu')
             # checkpoint state dict
             checkpoint_state_dict = checkpoint.pop("model")
-            model.load_state_dict(checkpoint_state_dict)
+            # check
+            new_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: