From b96855e3529819b6b0bb48ed4b744e42ade780b6 Mon Sep 17 00:00:00 2001 From: Zhiqiang Wang Date: Fri, 15 Oct 2021 19:57:19 +0800 Subject: [PATCH] Fixing ImprotError --- test/test_models_utils.py | 83 ---------------------- test/test_utils.py | 102 +++++++++++++++++++++++++-- yolort/models/__init__.py | 2 +- yolort/models/_utils.py | 65 ----------------- yolort/models/yolo.py | 2 +- yolort/models/yolo_module.py | 3 +- yolort/utils/__init__.py | 5 +- yolort/utils/update_module_state.py | 66 ++++++++++++++++- yolort/{ => v5}/utils/activations.py | 8 ++- 9 files changed, 172 insertions(+), 164 deletions(-) delete mode 100644 test/test_models_utils.py rename yolort/{ => v5}/utils/activations.py (74%) diff --git a/test/test_models_utils.py b/test/test_models_utils.py deleted file mode 100644 index 34d7b660..00000000 --- a/test/test_models_utils.py +++ /dev/null @@ -1,83 +0,0 @@ -# Copyright (c) 2021, Zhiqiang Wang. All Rights Reserved. -import pytest -import torch -from torch import Tensor -from yolort import models -from yolort.models import YOLOv5 -from yolort.models._utils import load_from_ultralytics - - -@pytest.mark.parametrize( - "arch, version, upstream_version, hash_prefix, use_p6", - [ - ("yolov5s", "r4.0", "v4.0", "9ca9a642", False), - ("yolov5s", "r4.0", "v6.0", "c3b140f3", False), - ], -) -def test_load_from_ultralytics( - arch: str, - version: str, - upstream_version: str, - hash_prefix: str, - use_p6: bool, -): - checkpoint_path = f"{arch}_{version}_{hash_prefix}" - base_url = "https://github.com/ultralytics/yolov5/releases/download/" - model_url = f"{base_url}/{upstream_version}/{arch}.pt" - - torch.hub.download_url_to_file( - model_url, - checkpoint_path, - hash_prefix=hash_prefix, - ) - model_info = load_from_ultralytics(checkpoint_path, version=version) - assert isinstance(model_info, dict) - assert model_info["num_classes"] == 80 - assert model_info["size"] == arch.replace("yolov5", "") - assert model_info["use_p6"] == use_p6 - assert len(model_info["strides"]) == 4 if use_p6 else 3 - - -@pytest.mark.parametrize( - "arch, version, upstream_version, hash_prefix", - [("yolov5s", "r4.0", "v4.0", "9ca9a642")], -) -def test_load_from_yolov5( - arch: str, - version: str, - upstream_version: str, - hash_prefix: str, -): - img_path = "test/assets/bus.jpg" - checkpoint_path = f"{arch}_{version}_{hash_prefix}" - - base_url = "https://github.com/ultralytics/yolov5/releases/download/" - model_url = f"{base_url}/{upstream_version}/{arch}.pt" - - torch.hub.download_url_to_file( - model_url, - checkpoint_path, - hash_prefix=hash_prefix, - ) - - model_yolov5 = YOLOv5.load_from_yolov5(checkpoint_path, 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.eval() - out = model.predict(img_path) - - torch.testing.assert_close( - out_from_yolov5[0]["scores"], out[0]["scores"], rtol=0, atol=0 - ) - torch.testing.assert_close( - out_from_yolov5[0]["labels"], out[0]["labels"], rtol=0, atol=0 - ) - torch.testing.assert_close( - out_from_yolov5[0]["boxes"], out[0]["boxes"], rtol=0, atol=0 - ) diff --git a/test/test_utils.py b/test/test_utils.py index 60f568b9..200beea6 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -3,16 +3,94 @@ import pytest import torch from torch import Tensor -from yolort.models import yolov5s +from yolort import models +from yolort.models import YOLOv5 from yolort.utils import ( FeatureExtractor, get_image_from_url, + load_from_ultralytics, read_image_to_tensor, ) from yolort.utils.image_utils import box_cxcywh_to_xyxy from yolort.v5 import letterbox, scale_coords +@pytest.mark.parametrize( + "arch, version, upstream_version, hash_prefix, use_p6", + [ + ("yolov5s", "r4.0", "v4.0", "9ca9a642", False), + ("yolov5s", "r4.0", "v6.0", "c3b140f3", False), + ], +) +def test_load_from_ultralytics( + arch: str, + version: str, + upstream_version: str, + hash_prefix: str, + use_p6: bool, +): + checkpoint_path = f"{arch}_{version}_{hash_prefix}" + base_url = "https://github.com/ultralytics/yolov5/releases/download/" + model_url = f"{base_url}/{upstream_version}/{arch}.pt" + + torch.hub.download_url_to_file( + model_url, + checkpoint_path, + hash_prefix=hash_prefix, + ) + model_info = load_from_ultralytics(checkpoint_path, version=version) + assert isinstance(model_info, dict) + assert model_info["num_classes"] == 80 + assert model_info["size"] == arch.replace("yolov5", "") + assert model_info["use_p6"] == use_p6 + assert len(model_info["strides"]) == 4 if use_p6 else 3 + + +@pytest.mark.parametrize( + "arch, version, upstream_version, hash_prefix", + [("yolov5s", "r4.0", "v4.0", "9ca9a642")], +) +def test_load_from_yolov5( + arch: str, + version: str, + upstream_version: str, + hash_prefix: str, +): + img_path = "test/assets/bus.jpg" + checkpoint_path = f"{arch}_{version}_{hash_prefix}" + + base_url = "https://github.com/ultralytics/yolov5/releases/download/" + model_url = f"{base_url}/{upstream_version}/{arch}.pt" + + torch.hub.download_url_to_file( + model_url, + checkpoint_path, + hash_prefix=hash_prefix, + ) + + model_yolov5 = YOLOv5.load_from_yolov5(checkpoint_path, 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.eval() + out = model.predict(img_path) + + torch.testing.assert_close( + out_from_yolov5[0]["scores"], out[0]["scores"], rtol=0, atol=0 + ) + torch.testing.assert_close( + out_from_yolov5[0]["labels"], out[0]["labels"], rtol=0, atol=0 + ) + torch.testing.assert_close( + out_from_yolov5[0]["boxes"], out[0]["boxes"], rtol=0, atol=0 + ) + + def test_read_image_to_tensor(): N, H, W = 3, 720, 360 img = np.random.randint(0, 255, (H, W, N), dtype="uint8") # As a dummy image @@ -76,21 +154,31 @@ def test_scale_coords(): torch.testing.assert_close(box_coords_scaled, exp_coords) -@pytest.mark.parametrize("b, h, w", [(8, 640, 640), (4, 416, 320), (8, 320, 416)]) -def test_feature_extractor(b, h, w): +@pytest.mark.parametrize( + "batch_size, height, width", + [ + (8, 640, 640), + (4, 416, 320), + (8, 320, 416), + ] +) +@pytest.mark.parametrize("arch", ["yolov5n", "yolov5s"]) +def test_feature_extractor(batch_size, height, width, arch): c = 3 in_channels = [128, 256, 512] strides = [8, 16, 32] num_outputs = 85 expected_features = [ - (b, inc, h // s, w // s) for inc, s in zip(in_channels, strides) + (batch_size, inc, height // s, width // s) for inc, s in zip(in_channels, strides) + ] + expected_head_outputs = [ + (batch_size, c, height // s, width // s, num_outputs) for s in strides ] - expected_head_outputs = [(b, c, h // s, w // s, num_outputs) for s in strides] - model = yolov5s() + model = models.__dict__[arch]() model = model.train() yolo_features = FeatureExtractor(model.model, return_layers=["backbone", "head"]) - images = torch.rand(b, c, h, w) + images = torch.rand(batch_size, c, height, width) targets = torch.rand(61, 6) intermediate_features = yolo_features(images, targets) features = intermediate_features["backbone"] diff --git a/yolort/models/__init__.py b/yolort/models/__init__.py index 26021d44..c9c5fb0c 100644 --- a/yolort/models/__init__.py +++ b/yolort/models/__init__.py @@ -3,8 +3,8 @@ from torch import nn -from yolort.utils.activations import Hardswish, SiLU from yolort.v5 import Conv +from yolort.v5.utils.activations import Hardswish, SiLU from .yolo import YOLO from .yolo_module import YOLOv5 diff --git a/yolort/models/_utils.py b/yolort/models/_utils.py index 3ff332cc..37e129d9 100644 --- a/yolort/models/_utils.py +++ b/yolort/models/_utils.py @@ -5,71 +5,6 @@ from torch import nn, Tensor from torchvision.ops import box_convert, box_iou -from yolort.utils import ModuleStateUpdate -from yolort.v5 import load_yolov5_model, get_yolov5_size - - -def load_from_ultralytics(checkpoint_path: str, version: str = "r6.0"): - """ - Allows the user to load model state file from the checkpoint trained from - the ultralytics/yolov5. - - Args: - checkpoint_path (str): Path of the YOLOv5 checkpoint model. - version (str): upstream version released by the ultralytics/yolov5, Possible - values are ["r3.1", "r4.0", "r6.0"]. Default: "r6.0". - """ - - assert version in [ - "r3.1", - "r4.0", - "r6.0", - ], "Currently does not support this version." - - checkpoint_yolov5 = load_yolov5_model(checkpoint_path) - num_classes = checkpoint_yolov5.yaml["nc"] - strides = checkpoint_yolov5.stride - anchor_grids = checkpoint_yolov5.yaml["anchors"] - depth_multiple = checkpoint_yolov5.yaml["depth_multiple"] - width_multiple = checkpoint_yolov5.yaml["width_multiple"] - - use_p6 = False - if len(strides) == 4: - use_p6 = True - - if use_p6: - inner_block_maps = {"0": "9", "1": "10", "3": "13", "4": "14"} - layer_block_maps = {"0": "17", "1": "18", "2": "20", "3": "21", "4": "23"} - else: - inner_block_maps = {"0": "9", "1": "10", "3": "13", "4": "14"} - layer_block_maps = {"0": "17", "1": "18", "2": "20", "3": "21", "4": "23"} - - module_state_updater = ModuleStateUpdate( - arch=None, - depth_multiple=depth_multiple, - width_multiple=width_multiple, - version=version, - num_classes=num_classes, - inner_block_maps=inner_block_maps, - layer_block_maps=layer_block_maps, - use_p6=use_p6, - ) - module_state_updater.updating(checkpoint_yolov5) - state_dict = module_state_updater.model.state_dict() - - size = get_yolov5_size(depth_multiple, width_multiple) - - return { - "num_classes": num_classes, - "depth_multiple": depth_multiple, - "width_multiple": width_multiple, - "strides": strides, - "anchor_grids": anchor_grids, - "use_p6": use_p6, - "size": size, - "state_dict": state_dict, - } - def _evaluate_iou(target, pred): """ diff --git a/yolort/models/yolo.py b/yolort/models/yolo.py index 31046905..6095722c 100644 --- a/yolort/models/yolo.py +++ b/yolort/models/yolo.py @@ -7,7 +7,7 @@ from torch import nn, Tensor from torchvision.models.utils import load_state_dict_from_url -from ._utils import load_from_ultralytics +from yolort.utils import load_from_ultralytics from .anchor_utils import AnchorGenerator from .backbone_utils import darknet_pan_backbone from .box_head import YOLOHead, SetCriterion, PostProcess diff --git a/yolort/models/yolo_module.py b/yolort/models/yolo_module.py index 9ddeda97..40562ce4 100644 --- a/yolort/models/yolo_module.py +++ b/yolort/models/yolo_module.py @@ -10,8 +10,9 @@ from torchvision.io import read_image from yolort.data import COCOEvaluator, contains_any_tensor +from yolort.utils import load_from_ultralytics from . import yolo -from ._utils import _evaluate_iou, load_from_ultralytics +from ._utils import _evaluate_iou from .transform import YOLOTransform __all__ = ["YOLOv5"] diff --git a/yolort/utils/__init__.py b/yolort/utils/__init__.py index 3b2b19f8..d2185361 100644 --- a/yolort/utils/__init__.py +++ b/yolort/utils/__init__.py @@ -2,14 +2,15 @@ from .hooks import FeatureExtractor from .image_utils import cv2_imshow, get_image_from_url, read_image_to_tensor -from .update_module_state import ModuleStateUpdate +from .update_module_state import load_from_ultralytics + __all__ = [ "FeatureExtractor", - "ModuleStateUpdate", "cv2_imshow", "get_image_from_url", "get_callable_dict", + "load_from_ultralytics", "read_image_to_tensor", ] diff --git a/yolort/utils/update_module_state.py b/yolort/utils/update_module_state.py index 34214bab..e65f1b8b 100644 --- a/yolort/utils/update_module_state.py +++ b/yolort/utils/update_module_state.py @@ -4,8 +4,69 @@ from torch import nn +from yolort.v5 import load_yolov5_model, get_yolov5_size from yolort.models import yolo -from yolort.v5 import get_yolov5_size + + +def load_from_ultralytics(checkpoint_path: str, version: str = "r6.0"): + """ + Allows the user to load model state file from the checkpoint trained from + the ultralytics/yolov5. + + Args: + checkpoint_path (str): Path of the YOLOv5 checkpoint model. + version (str): upstream version released by the ultralytics/yolov5, Possible + values are ["r3.1", "r4.0", "r6.0"]. Default: "r6.0". + """ + + assert version in [ + "r3.1", + "r4.0", + "r6.0", + ], "Currently does not support this version." + + checkpoint_yolov5 = load_yolov5_model(checkpoint_path) + num_classes = checkpoint_yolov5.yaml["nc"] + strides = checkpoint_yolov5.stride + anchor_grids = checkpoint_yolov5.yaml["anchors"] + depth_multiple = checkpoint_yolov5.yaml["depth_multiple"] + width_multiple = checkpoint_yolov5.yaml["width_multiple"] + + use_p6 = False + if len(strides) == 4: + use_p6 = True + + if use_p6: + inner_block_maps = {"0": "9", "1": "10", "3": "13", "4": "14"} + layer_block_maps = {"0": "17", "1": "18", "2": "20", "3": "21", "4": "23"} + else: + inner_block_maps = {"0": "9", "1": "10", "3": "13", "4": "14"} + layer_block_maps = {"0": "17", "1": "18", "2": "20", "3": "21", "4": "23"} + + module_state_updater = ModuleStateUpdate( + depth_multiple=depth_multiple, + width_multiple=width_multiple, + version=version, + num_classes=num_classes, + inner_block_maps=inner_block_maps, + layer_block_maps=layer_block_maps, + use_p6=use_p6, + ) + module_state_updater.updating(checkpoint_yolov5) + state_dict = module_state_updater.model.state_dict() + + size = get_yolov5_size(depth_multiple, width_multiple) + + return { + "num_classes": num_classes, + "depth_multiple": depth_multiple, + "width_multiple": width_multiple, + "strides": strides, + "anchor_grids": anchor_grids, + "use_p6": use_p6, + "size": size, + "state_dict": state_dict, + } class ModuleStateUpdate: @@ -56,7 +117,7 @@ def __init__( weights_name = ( f"yolov5_darknet_pan_{yolov5_size}_{version.replace('.', '')}_coco" ) - model = yolo.build_model( + self.model = yolo.build_model( backbone_name, depth_multiple, width_multiple, @@ -65,7 +126,6 @@ def __init__( num_classes=num_classes, use_p6=use_p6, ) - self.model = model def updating(self, state_dict): # Obtain module state diff --git a/yolort/utils/activations.py b/yolort/v5/utils/activations.py similarity index 74% rename from yolort/utils/activations.py rename to yolort/v5/utils/activations.py index 8dec1dc1..1d14d945 100644 --- a/yolort/utils/activations.py +++ b/yolort/v5/utils/activations.py @@ -1,3 +1,8 @@ +# YOLOv5 🚀 by Ultralytics, GPL-3.0 license +""" +Activation functions +""" + import torch import torch.nn as nn import torch.nn.functional as F @@ -26,4 +31,5 @@ class Hardswish(nn.Module): @staticmethod def forward(x): - return x * F.hardtanh(x + 3, 0.0, 6.0) / 6.0 + # return x * F.hardsigmoid(x) # for torchscript and CoreML + return x * F.hardtanh(x + 3, 0., 6.) / 6. # for torchscript, CoreML and ONNX