From 725f6ebd7659f90962649bc8d2afeac532a7e7f9 Mon Sep 17 00:00:00 2001 From: Zhiqiang Wang Date: Sun, 6 Mar 2022 11:34:09 +0800 Subject: [PATCH 1/4] Use C3 following the model specification of YOLOv5 --- yolort/models/box_head.py | 2 +- yolort/models/yolo_lite.py | 22 +++++++++++++++++++--- 2 files changed, 20 insertions(+), 4 deletions(-) diff --git a/yolort/models/box_head.py b/yolort/models/box_head.py index 4eb6f73b..4210f8ff 100644 --- a/yolort/models/box_head.py +++ b/yolort/models/box_head.py @@ -20,7 +20,7 @@ def __init__( ): super().__init__() if not isinstance(in_channels, list): - in_channels = [in_channels] * num_anchors + in_channels = [in_channels] * len(strides) self.num_anchors = num_anchors # anchors self.num_classes = num_classes self.num_outputs = num_classes + 5 # number of outputs per anchor diff --git a/yolort/models/yolo_lite.py b/yolort/models/yolo_lite.py index 0d3b15bf..430da93a 100644 --- a/yolort/models/yolo_lite.py +++ b/yolort/models/yolo_lite.py @@ -5,6 +5,8 @@ from torchvision.ops import misc as misc_nn_ops from torchvision.ops.feature_pyramid_network import FeaturePyramidNetwork, LastLevelMaxPool +from .anchor_utils import AnchorGenerator +from .box_head import YOLOHead from .yolo import YOLO __all__ = ["yolov5_mobilenet_v3_small_fpn"] @@ -79,7 +81,7 @@ def mobilenet_backbone( out_channels = 256 if returned_layers is None: - returned_layers = [num_stages - 2, num_stages - 1] + returned_layers = [num_stages - 3, num_stages - 2, num_stages - 1] assert min(returned_layers) >= 0 and max(returned_layers) < num_stages return_layers = {f"{stage_indices[k]}": str(v) for v, k in enumerate(returned_layers)} @@ -113,8 +115,22 @@ def _yolov5_mobilenet_v3_small_fpn( pretrained_backbone, trainable_layers=trainable_backbone_layers, ) - - model = YOLO(backbone, num_classes, **kwargs) + strides = [8, 16, 32, 64] + anchor_grids = [ + [19, 27, 44, 40, 38, 94], + [96, 68, 86, 152, 180, 137], + [140, 301, 303, 264, 238, 542], + [436, 615, 739, 380, 925, 792], + ] + anchor_generator = AnchorGenerator(strides, anchor_grids) + + head = YOLOHead( + backbone.out_channels, + anchor_generator.num_anchors, + anchor_generator.strides, + num_classes, + ) + model = YOLO(backbone, num_classes, anchor_generator=anchor_generator, head=head, **kwargs) return model From cb1e293701e02c7b5a9f9f0201261ead8ec83871 Mon Sep 17 00:00:00 2001 From: Zhiqiang Wang Date: Sun, 6 Mar 2022 11:48:29 +0800 Subject: [PATCH 2/4] Add unittest for yolov5_mobilenet_v3_small_fpn --- test/test_models.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/test/test_models.py b/test/test_models.py index 1b55ecb2..0ef183d5 100644 --- a/test/test_models.py +++ b/test/test_models.py @@ -12,6 +12,7 @@ from yolort.models.backbone_utils import darknet_pan_backbone from yolort.models.box_head import YOLOHead, PostProcess, SetCriterion from yolort.models.transformer import darknet_tan_backbone +from yolort.models.yolo_lite import yolov5_mobilenet_v3_small_fpn from yolort.v5 import get_yolov5_size, attempt_download @@ -420,3 +421,18 @@ def test_load_from_yolov5_torchscript(arch, size_divisible, version, upstream_ve torch.testing.assert_close(out[0]["scores"], out_script[1][0]["scores"], rtol=0, atol=0) torch.testing.assert_close(out[0]["labels"], out_script[1][0]["labels"], rtol=0, atol=0) torch.testing.assert_close(out[0]["boxes"], out_script[1][0]["boxes"], rtol=0, atol=0) + + +def test_yolov5_mobilenet_v3_small_fpn(): + + model = yolov5_mobilenet_v3_small_fpn() + model = model.eval() + + images = torch.rand(4, 3, 320, 320) + out = model(images) + assert isinstance(out, list) + assert len(out) == 4 + assert isinstance(out[0], dict) + assert isinstance(out[0]["boxes"], Tensor) + assert isinstance(out[0]["labels"], Tensor) + assert isinstance(out[0]["scores"], Tensor) From 6ee9f36389637e6e91bd162b6d1f69f98198c7ff Mon Sep 17 00:00:00 2001 From: Zhiqiang Wang Date: Sun, 6 Mar 2022 12:11:39 +0800 Subject: [PATCH 3/4] Fix format --- yolort/models/yolo_lite.py | 60 +++++++++++++++++++++++++------------- 1 file changed, 39 insertions(+), 21 deletions(-) diff --git a/yolort/models/yolo_lite.py b/yolort/models/yolo_lite.py index 430da93a..4ad5bb06 100644 --- a/yolort/models/yolo_lite.py +++ b/yolort/models/yolo_lite.py @@ -1,9 +1,12 @@ +from typing import Dict, List, Callable, Optional + from torch import nn from torchvision.models import mobilenet from torchvision.models._utils import IntermediateLayerGetter from torchvision.models.detection.backbone_utils import _validate_trainable_layers from torchvision.ops import misc as misc_nn_ops -from torchvision.ops.feature_pyramid_network import FeaturePyramidNetwork, LastLevelMaxPool +from torchvision.ops.feature_pyramid_network import ExtraFPNBlock, FeaturePyramidNetwork, LastLevelMaxPool +from yolort.utils import load_state_dict_from_url from .anchor_utils import AnchorGenerator from .box_head import YOLOHead @@ -33,7 +36,14 @@ class BackboneWithFPN(nn.Module): out_channels (int): the number of channels in the FPN """ - def __init__(self, backbone, return_layers, in_channels_list, out_channels, extra_blocks=None): + def __init__( + self, + backbone: nn.Module, + return_layers: Dict[str, str], + in_channels_list: List[int], + out_channels: int, + extra_blocks: Optional[ExtraFPNBlock] = None, + ) -> None: super().__init__() if extra_blocks is None: @@ -55,12 +65,12 @@ def forward(self, x): def mobilenet_backbone( - backbone_name, - pretrained, - norm_layer=misc_nn_ops.FrozenBatchNorm2d, - trainable_layers=2, - returned_layers=None, -): + backbone_name: str, + pretrained: bool, + norm_layer: Callable[..., nn.Module] = misc_nn_ops.FrozenBatchNorm2d, + trainable_layers: int = 2, + returned_layers: Optional[List[int]] = None, +) -> nn.Module: backbone = mobilenet.__dict__[backbone_name](pretrained=pretrained, norm_layer=norm_layer).features # Gather the indices of blocks which are strided. These are the locations of C1, ..., Cn-1 blocks. @@ -96,12 +106,12 @@ def mobilenet_backbone( def _yolov5_mobilenet_v3_small_fpn( - weights_name, - pretrained=False, - progress=True, - num_classes=80, - pretrained_backbone=True, - trainable_backbone_layers=None, + weights_name: str, + pretrained: bool = False, + progress: bool = True, + num_classes: int = 80, + pretrained_backbone: bool = True, + trainable_backbone_layers: Optional[int] = None, **kwargs, ): trainable_backbone_layers = _validate_trainable_layers( @@ -123,7 +133,6 @@ def _yolov5_mobilenet_v3_small_fpn( [436, 615, 739, 380, 925, 792], ] anchor_generator = AnchorGenerator(strides, anchor_grids) - head = YOLOHead( backbone.out_channels, anchor_generator.num_anchors, @@ -131,16 +140,25 @@ def _yolov5_mobilenet_v3_small_fpn( num_classes, ) model = YOLO(backbone, num_classes, anchor_generator=anchor_generator, head=head, **kwargs) - + if pretrained: + if model_urls.get(weights_name, None) is None: + raise ValueError(f"No checkpoint is available for model {weights_name}") + state_dict = load_state_dict_from_url(model_urls["retinanet_resnet50_fpn_coco"], progress=progress) + model.load_state_dict(state_dict) return model +model_urls = { + "yolov5_mobilenet_v3_small_fpn_coco": None, +} + + def yolov5_mobilenet_v3_small_fpn( - pretrained=False, - progress=True, - num_classes=80, - pretrained_backbone=True, - trainable_backbone_layers=None, + pretrained: bool = False, + progress: bool = True, + num_classes: int = 80, + pretrained_backbone: bool = True, + trainable_backbone_layers: Optional[int] = None, **kwargs, ): """ From 39b4342c84ae5620aa244ed307b960fc1e23ea26 Mon Sep 17 00:00:00 2001 From: Zhiqiang Wang Date: Sun, 6 Mar 2022 12:16:20 +0800 Subject: [PATCH 4/4] Minor fix for weights_name --- yolort/models/yolo_lite.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/yolort/models/yolo_lite.py b/yolort/models/yolo_lite.py index 4ad5bb06..627f0863 100644 --- a/yolort/models/yolo_lite.py +++ b/yolort/models/yolo_lite.py @@ -143,7 +143,7 @@ def _yolov5_mobilenet_v3_small_fpn( if pretrained: if model_urls.get(weights_name, None) is None: raise ValueError(f"No checkpoint is available for model {weights_name}") - state_dict = load_state_dict_from_url(model_urls["retinanet_resnet50_fpn_coco"], progress=progress) + state_dict = load_state_dict_from_url(model_urls[weights_name], progress=progress) model.load_state_dict(state_dict) return model