Add checkpoint and unit-test for yolov5m and yolov5l

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

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/test/test_models.py b/test/test_models.py
@@ -360,6 +360,9 @@ def test_torchscript(arch):
     [
         ("yolov5s", "r4.0", "v4.0", "9ca9a642"),
         ("yolov5n", "r6.0", "v6.0", "649e089f"),
+        ("yolov5s", "r6.0", "v6.0", "c3b140f3"),
+        ("yolov5m", "r6.0", "v6.0", "4947bf56"),
+        ("yolov5l", "r6.0", "v6.0", "84f0a2d6"),
     ],
 )
 def test_load_from_yolov5(

diff --git a/yolort/models/yolo.py b/yolort/models/yolo.py
@@ -292,9 +292,9 @@ def build_model(
     "yolov5_darknet_pan_n6_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_n6_r60_coco.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_m_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_m_r60_coco.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_l_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_l_r60_coco.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",
 }