Add checkpoint and unit-test for yolov5n v6.0

zhiqwang · Oct 23, 2021 · 4af9dc5 · 4af9dc5
1 parent 46cb4aa
commit 4af9dc5
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 5 deletions.
diff --git a/test/test_models.py b/test/test_models.py
@@ -357,7 +357,10 @@ def test_torchscript(arch):
 
 @pytest.mark.parametrize(
     "arch, version, upstream_version, hash_prefix",
-    [("yolov5s", "r4.0", "v4.0", "9ca9a642")],
+    [
+        ("yolov5s", "r4.0", "v4.0", "9ca9a642"),
+        ("yolov5n", "r6.0", "v6.0", "649e089f"),
+    ],
 )
 def test_load_from_yolov5(
     arch: str,
@@ -376,16 +379,25 @@ def test_load_from_yolov5(
         checkpoint_path,
         hash_prefix=hash_prefix,
     )
+    score_thresh = 0.25
 
-    model_yolov5 = YOLOv5.load_from_yolov5(checkpoint_path, version=version)
+    model_yolov5 = YOLOv5.load_from_yolov5(
+        checkpoint_path,
+        score_thresh=score_thresh,
+        version=version,
+    )
     model_yolov5.eval()
     out_from_yolov5 = model_yolov5.predict(img_path)
     assert isinstance(out_from_yolov5[0], dict)
     assert isinstance(out_from_yolov5[0]["boxes"], Tensor)
     assert isinstance(out_from_yolov5[0]["labels"], Tensor)
     assert isinstance(out_from_yolov5[0]["scores"], Tensor)
 
-    model = models.__dict__[arch](pretrained=True, score_thresh=0.25)
+    model = models.__dict__[arch](
+        upstream_version=version,
+        pretrained=True,
+        score_thresh=score_thresh,
+    )
     model.eval()
     out = model.predict(img_path)
 

diff --git a/yolort/models/yolo.py b/yolort/models/yolo.py
@@ -288,9 +288,9 @@ def build_model(
     "yolov5_darknet_pan_m_r40_coco": f"{model_urls_root_r30}/yolov5_darknet_pan_m_r40_coco-d295cb02.pt",
     "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.pt",
+    "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_s_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_s_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_m6_r60_coco": f"{model_urls_root_r52}/yolov5_darknet_pan_m6_r60_coco.pt",