From 4af9dc5f0509c00081739871f3d377eaba9784e5 Mon Sep 17 00:00:00 2001 From: zhiqiang Date: Sun, 24 Oct 2021 07:55:48 +0800 Subject: [PATCH] Add checkpoint and unit-test for yolov5n v6.0 --- test/test_models.py | 18 +++++++++++++++--- yolort/models/yolo.py | 4 ++-- 2 files changed, 17 insertions(+), 5 deletions(-) diff --git a/test/test_models.py b/test/test_models.py index 43d72856..b2845706 100644 --- 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,8 +379,13 @@ 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) @@ -385,7 +393,11 @@ def test_load_from_yolov5( 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 index 7619e687..cf9a93a2 100644 --- 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",