Add checkpoint and unit-test for yolov5 P6 releases

zhiqwang · Oct 24, 2021 · 1a2ac75 · 1a2ac75
1 parent 2ff04f0
commit 1a2ac75
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 6 deletions.
diff --git a/test/test_models.py b/test/test_models.py
@@ -363,6 +363,9 @@ def test_torchscript(arch):
         ("yolov5s", "r6.0", "v6.0", "c3b140f3"),
         ("yolov5m", "r6.0", "v6.0", "4947bf56"),
         ("yolov5l", "r6.0", "v6.0", "84f0a2d6"),
+        ("yolov5n6", "r6.0", "v6.0", "beecbbae"),
+        ("yolov5s6", "r6.0", "v6.0", "4fe88758"),
+        ("yolov5m6", "r6.0", "v6.0", "ae177092"),
     ],
 )
 def test_load_from_yolov5(

diff --git a/yolort/models/__init__.py b/yolort/models/__init__.py
@@ -8,7 +8,10 @@
 from .yolo import YOLO
 from .yolo_module import YOLOv5
 
-__all__ = ["YOLO", "YOLOv5", "yolov5n", "yolov5s", "yolov5m", "yolov5l", "yolov5ts"]
+__all__ = [
+    "YOLO", "YOLOv5", "yolov5n", "yolov5n6", "yolov5s", "yolov5s6",
+    "yolov5m", "yolov5m6", "yolov5l", "yolov5ts"
+]
 
 
 def yolov5n(

diff --git a/yolort/models/yolo.py b/yolort/models/yolo.py
@@ -289,11 +289,11 @@ def build_model(
     "yolov5_darknet_pan_l_r40_coco": f"{model_urls_root_r30}/yolov5_darknet_pan_l_r40_coco-4416841f.pt",
     # Path Aggregation Network 6.0
     "yolov5_darknet_pan_n_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_n_r60_coco-bc15659e.pt",
-    "yolov5_darknet_pan_n6_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_n6_r60_coco.pt",
+    "yolov5_darknet_pan_n6_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_n6_r60_coco-4e823e0f.pt",
     "yolov5_darknet_pan_s_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_s_r60_coco-9f44bf3f.pt",
-    "yolov5_darknet_pan_s6_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_s6_r60_coco.pt",
+    "yolov5_darknet_pan_s6_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_s6_r60_coco-b4ff1fc2.pt",
     "yolov5_darknet_pan_m_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_m_r60_coco-58d32352.pt",
-    "yolov5_darknet_pan_m6_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_m6_r60_coco.pt",
+    "yolov5_darknet_pan_m6_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_m6_r60_coco-cc010533.pt",
     "yolov5_darknet_pan_l_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_l_r60_coco-321d8dcd.pt",
     # Tranformer Attention Network
     "yolov5_darknet_tan_s_r40_coco": f"{model_urls_root_r30}/yolov5_darknet_tan_s_r40_coco-fe1069ce.pt",

diff --git a/yolort/utils/update_module_state.py b/yolort/utils/update_module_state.py
@@ -32,8 +32,9 @@ def convert_yolov5_to_yolort(
     model_state_dict = model_info["state_dict"]
 
     size = model_info["size"]
-    output_path = output_path / f"{prefix}_{size}_{version.replace('.', '')}_{postfix}"
-    torch.save(model_state_dict, output_path)
+    use_p6 = "6" if model_info["use_p6"] else ""
+    output_postfix = f"{prefix}_{size}{use_p6}_{version.replace('.', '')}_{postfix}"
+    torch.save(model_state_dict, output_path / output_postfix)
 
 
 def load_from_ultralytics(