Resolve version conflicts

hubconf.py

@@ -1,4 +1,4 @@
 # Optional list of dependencies required by the package
 dependencies = ["torch", "torchvision"]
 
-from yolort.models import yolov5s, yolov5m, yolov5l
+from yolort.models import yolov5s, yolov5m, yolov5l, yolov5ts

test/test_models.py

@@ -131,10 +131,10 @@ def _init_test_backbone_with_pan_r3_1(self):
         backbone_name = "darknet_s_r3_1"
         depth_multiple = 0.33
         width_multiple = 0.5
-        backbone_with_fpn = darknet_pan_backbone(
+        backbone_with_pan = darknet_pan_backbone(
             backbone_name, depth_multiple, width_multiple
         )
-        return backbone_with_fpn
+        return backbone_with_pan
 
     def test_backbone_with_pan_r3_1(self):
         N, H, W = 4, 416, 352
@@ -154,10 +154,10 @@ def _init_test_backbone_with_pan_r4_0(self):
         backbone_name = "darknet_s_r4_0"
         depth_multiple = 0.33
         width_multiple = 0.5
-        backbone_with_fpn = darknet_pan_backbone(
+        backbone_with_pan = darknet_pan_backbone(
             backbone_name, depth_multiple, width_multiple
         )
-        return backbone_with_fpn
+        return backbone_with_pan
 
     def test_backbone_with_pan_r4_0(self):
         N, H, W = 4, 416, 352
@@ -173,21 +173,21 @@ def test_backbone_with_pan_r4_0(self):
         assert tuple(out[2].shape) == (N, *out_shape[2])
         _check_jit_scriptable(model, (x,))
 
-    def _init_test_backbone_with_pan_tr(self):
+    def _init_test_backbone_with_tan_r4_0(self):
         backbone_name = "darknet_s_r4_0"
         depth_multiple = 0.33
         width_multiple = 0.5
-        backbone_with_fpn_tr = darknet_tan_backbone(
+        backbone_with_tan = darknet_tan_backbone(
             backbone_name, depth_multiple, width_multiple
         )
-        return backbone_with_fpn_tr
+        return backbone_with_tan
 
-    def test_backbone_with_pan_tr(self):
+    def test_backbone_with_tan_r4_0(self):
         N, H, W = 4, 416, 352
         out_shape = self._get_feature_shapes(H, W)
 
         x = torch.rand(N, 3, H, W)
-        model = self._init_test_backbone_with_pan_tr()
+        model = self._init_test_backbone_with_tan_r4_0()
         out = model(x)
 
         assert len(out) == 3
@@ -278,7 +278,7 @@ def test_criterion(self):
         assert isinstance(losses["objectness"], Tensor)
 
 
-@pytest.mark.parametrize("arch", ["yolov5s", "yolov5m", "yolov5l", "yolotr"])
+@pytest.mark.parametrize("arch", ["yolov5s", "yolov5m", "yolov5l", "yolov5ts"])
 def test_torchscript(arch):
     model = models.__dict__[arch](pretrained=True, size=(320, 320), score_thresh=0.45)
     model.eval()
@@ -303,21 +303,21 @@ def test_torchscript(arch):
 
 
 @pytest.mark.parametrize(
-    "arch, version, hash_prefix", [("yolov5s", "v4.0", "9ca9a642")]
+    "arch, up_version, hash_prefix", [("yolov5s", "v4.0", "9ca9a642")]
 )
-def test_load_from_yolov5(arch, version, hash_prefix):
+def test_load_from_yolov5(arch, up_version, hash_prefix):
     img_path = "test/assets/bus.jpg"
     yolov5s_r40_path = Path(f"{arch}.pt")
 
     if not yolov5s_r40_path.exists():
         torch.hub.download_url_to_file(
-            f"https://github.com/ultralytics/yolov5/releases/download/{version}/{arch}.pt",
+            f"https://github.com/ultralytics/yolov5/releases/download/{up_version}/{arch}.pt",
             yolov5s_r40_path,
             hash_prefix=hash_prefix,
         )
 
-    yolov5 = YOLOv5()
-    model_yolov5 = yolov5.load_from_yolov5(yolov5s_r40_path, score_thresh=0.25)
+    version = up_version.replace("v", "r")
+    model_yolov5 = YOLOv5.load_from_yolov5(yolov5s_r40_path, version=version)
     model_yolov5.eval()
     out_from_yolov5 = model_yolov5.predict(img_path)
     assert isinstance(out_from_yolov5[0], dict)

test/test_onnx.py

@@ -116,7 +116,7 @@ def get_test_images(self):
         [
             ("yolov5s", "r3.1"),
             ("yolov5m", "r4.0"),
-            # ('yolotr', 'r4.0'),
+            # ("yolov5ts", "r4.0"),
         ],
     )
     def test_yolort_export_onnx(self, arch, upstream_version):

yolort/models/__init__.py

@@ -8,7 +8,7 @@
 from .yolo import YOLO
 from .yolo_module import YOLOv5
 
-__all__ = ["YOLO", "YOLOv5", "yolov5s", "yolov5m", "yolov5l", "yolotr"]
+__all__ = ["YOLO", "YOLOv5", "yolov5s", "yolov5m", "yolov5l", "yolov5ts"]
 
 
 def yolov5s(
@@ -80,7 +80,7 @@ def yolov5l(
     return model
 
 
-def yolotr(
+def yolov5ts(
     upstream_version: str = "r4.0", export_friendly: bool = False, **kwargs: Any
 ):
     """

yolort/models/_utils.py

@@ -9,15 +9,21 @@
 from yolort.v5 import load_yolov5_model, get_yolov5_size
 
 
-def load_from_ultralytics(checkpoint_path: str, version: str = "r4.0"):
+def load_from_ultralytics(checkpoint_path: str, version: str = "r6.0"):
     """
     Load YOLOv5 state from the checkpoint trained from the ultralytics.
 
     Args:
         checkpoint_path (str): Path of the YOLOv5 checkpoint model.
-        version (str): upstream version released by the ultralytics/yolov5,
-            versions r3.1 and r4.0 are currently supported.
+        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 only supports version 'r3.1', 'r4.0' and 'r6.0'."
     checkpoint_yolov5 = load_yolov5_model(checkpoint_path)
     num_classes = checkpoint_yolov5.yaml["nc"]
     anchor_grids = checkpoint_yolov5.yaml["anchors"]

yolort/models/backbone_utils.py

@@ -24,7 +24,7 @@ class BackboneWithPAN(nn.Module):
         in_channels_list (List[int]): number of channels for each feature map
             that is returned, in the order they are present in the OrderedDict
         depth_multiple (float): depth multiplier
-        version (str): ultralytics release version: ["r3.1", "r4.0", "r6.0"]
+        version (str): Module version released by ultralytics: ["r3.1", "r4.0", "r6.0"].
 
     Attributes:
         out_channels (int): the number of channels in the PAN
@@ -55,7 +55,7 @@ def darknet_pan_backbone(
     width_multiple: float,
     pretrained: Optional[bool] = False,
     returned_layers: Optional[List[int]] = None,
-    version: str = "r4.0",
+    version: str = "r6.0",
 ):
     """
     Constructs a specified DarkNet backbone with PAN on top. Freezes the specified number of
@@ -64,7 +64,7 @@ def darknet_pan_backbone(
     Examples:
 
         >>> from models.backbone_utils import darknet_pan_backbone
-        >>> backbone = darknet_pan_backbone('darknet3_1', pretrained=True, trainable_layers=3)
+        >>> backbone = darknet_pan_backbone("darknet_s_r4_0")
         >>> # get some dummy image
         >>> x = torch.rand(1, 3, 64, 64)
         >>> # compute the output
@@ -75,15 +75,19 @@ def darknet_pan_backbone(
          ('2', torch.Size([1, 512, 2, 2]))]
 
     Args:
-        backbone_name (string): darknet architecture. Possible values are 'DarkNet', 'darknet_s_r3_1',
-           'darknet_m_r3_1', 'darknet_l_r3_1', 'darknet_s_r4_0', 'darknet_m_r4_0', 'darknet_l_r4_0'
-        norm_layer (torchvision.ops): it is recommended to use the default value. For details visit:
-            (https://github.com/facebookresearch/maskrcnn-benchmark/issues/267)
+        backbone_name (string): darknet architecture. Possible values are "darknet_s_r3_1",
+            "darknet_m_r3_1", "darknet_l_r3_1", "darknet_s_r4_0", "darknet_m_r4_0",
+            "darknet_l_r4_0", "darknet_s_r6_0", "darknet_m_r6_0", and "darknet_l_r6_0".
         pretrained (bool): If True, returns a model with backbone pre-trained on Imagenet
-        trainable_layers (int): number of trainable (not frozen) darknet layers starting from final block.
-            Valid values are between 0 and 5, with 5 meaning all backbone layers are trainable.
-        version (str): ultralytics release version: ["r3.1", "r4.0", "r6.0"]
+        version (str): Module version released by ultralytics. Possible values
+            are ["r3.1", "r4.0", "r6.0"]. Default: "r6.0".
     """
+    assert version in [
+        "r3.1",
+        "r4.0",
+        "r6.0",
+    ], "Currently only supports version 'r3.1', 'r4.0' and 'r6.0'."
+
     backbone = darknet.__dict__[backbone_name](pretrained=pretrained).features
 
     if returned_layers is None:

yolort/models/darknetv5.py

@@ -57,6 +57,10 @@ def __init__(
     ) -> None:
         super().__init__()
 
+        assert version in ["r3.1", "r4.0"], (
+            "Currently the module version used in DarkNetV5 is r3.1 or r4.0",
+        )
+
         if block is None:
             block = _block[version]
 
@@ -128,7 +132,7 @@ def forward(self, x: Tensor) -> Tensor:
 }
 
 
-def _darknet(
+def _darknetv5(
     arch: str, pretrained: bool, progress: bool, *args: Any, **kwargs: Any
 ) -> DarkNetV5:
     """
@@ -162,7 +166,9 @@ def darknet_s_r3_1(
         pretrained (bool): If True, returns a model pre-trained on ImageNet
         progress (bool): If True, displays a progress bar of the download to stderr
     """
-    return _darknet("darknet_s_r3.1", pretrained, progress, 0.33, 0.5, "r3.1", **kwargs)
+    return _darknetv5(
+        "darknet_s_r3.1", pretrained, progress, 0.33, 0.5, "r3.1", **kwargs
+    )
 
 
 def darknet_m_r3_1(
@@ -176,7 +182,7 @@ def darknet_m_r3_1(
         pretrained (bool): If True, returns a model pre-trained on ImageNet
         progress (bool): If True, displays a progress bar of the download to stderr
     """
-    return _darknet(
+    return _darknetv5(
         "darknet_m_r3.1", pretrained, progress, 0.67, 0.75, "r3.1", **kwargs
     )
 
@@ -192,7 +198,9 @@ def darknet_l_r3_1(
         pretrained (bool): If True, returns a model pre-trained on ImageNet
         progress (bool): If True, displays a progress bar of the download to stderr
     """
-    return _darknet("darknet_l_r3.1", pretrained, progress, 1.0, 1.0, "r3.1", **kwargs)
+    return _darknetv5(
+        "darknet_l_r3.1", pretrained, progress, 1.0, 1.0, "r3.1", **kwargs
+    )
 
 
 def darknet_s_r4_0(
@@ -206,7 +214,9 @@ def darknet_s_r4_0(
         pretrained (bool): If True, returns a model pre-trained on ImageNet
         progress (bool): If True, displays a progress bar of the download to stderr
     """
-    return _darknet("darknet_s_r4.0", pretrained, progress, 0.33, 0.5, "r4.0", **kwargs)
+    return _darknetv5(
+        "darknet_s_r4.0", pretrained, progress, 0.33, 0.5, "r4.0", **kwargs
+    )
 
 
 def darknet_m_r4_0(
@@ -220,7 +230,7 @@ def darknet_m_r4_0(
         pretrained (bool): If True, returns a model pre-trained on ImageNet
         progress (bool): If True, displays a progress bar of the download to stderr
     """
-    return _darknet(
+    return _darknetv5(
         "darknet_m_r4.0", pretrained, progress, 0.67, 0.75, "r4.0", **kwargs
     )
 
@@ -236,4 +246,6 @@ def darknet_l_r4_0(
         pretrained (bool): If True, returns a model pre-trained on ImageNet
         progress (bool): If True, displays a progress bar of the download to stderr
     """
-    return _darknet("darknet_l_r4.0", pretrained, progress, 1.0, 1.0, "r4.0", **kwargs)
+    return _darknetv5(
+        "darknet_l_r4.0", pretrained, progress, 1.0, 1.0, "r4.0", **kwargs
+    )

yolort/models/darknetv6.py

@@ -31,7 +31,7 @@ class DarkNetV6(nn.Module):
         depth_multiple (float): Depth multiplier
         width_multiple (float): Width multiplier - adjusts number of channels
             in each layer by this amount
-        version (str): ultralytics release version: r3.1 or r4.0
+        version (str): Module version released by ultralytics, set to r4.0.
         block: Module specifying inverted residual building block for darknet
         round_nearest (int): Round the number of channels in each layer to be
             a multiple of this number. Set to 1 to turn off rounding
@@ -42,7 +42,7 @@ def __init__(
         self,
         depth_multiple: float,
         width_multiple: float,
-        version: str,
+        version: str = "r4.0",
         block: Optional[Callable[..., nn.Module]] = None,
         stages_repeats: Optional[List[int]] = None,
         stages_out_channels: Optional[List[int]] = None,
@@ -51,6 +51,10 @@ def __init__(
     ) -> None:
         super().__init__()
 
+        assert version == "r4.0", (
+            "Currently the module version used in DarkNetV6 is r4.0",
+        )
+
         if block is None:
             block = C3
 
@@ -117,7 +121,7 @@ def forward(self, x: Tensor) -> Tensor:
         return self._forward_impl(x)
 
 
-def _darknet(
+def _darknetv6(
     arch: str, pretrained: bool, progress: bool, *args: Any, **kwargs: Any
 ) -> DarkNetV6:
     """
@@ -151,7 +155,7 @@ def darknet_s_r6_0(
         pretrained (bool): If True, returns a model pre-trained on ImageNet
         progress (bool): If True, displays a progress bar of the download to stderr
     """
-    return _darknet("darknet_s_r6.0", pretrained, progress, 0.33, 0.5, "r6.0", **kwargs)
+    return _darknetv6("darknet_s_r6.0", pretrained, progress, 0.33, 0.5, **kwargs)
 
 
 def darknet_m_r6_0(
@@ -165,9 +169,7 @@ def darknet_m_r6_0(
         pretrained (bool): If True, returns a model pre-trained on ImageNet
         progress (bool): If True, displays a progress bar of the download to stderr
     """
-    return _darknet(
-        "darknet_m_r6.0", pretrained, progress, 0.67, 0.75, "r6.0", **kwargs
-    )
+    return _darknetv6("darknet_m_r6.0", pretrained, progress, 0.67, 0.75, **kwargs)
 
 
 def darknet_l_r6_0(
@@ -181,4 +183,4 @@ def darknet_l_r6_0(
         pretrained (bool): If True, returns a model pre-trained on ImageNet
         progress (bool): If True, displays a progress bar of the download to stderr
     """
-    return _darknet("darknet_l_r6.0", pretrained, progress, 1.0, 1.0, "r6.0", **kwargs)
+    return _darknetv6("darknet_l_r6.0", pretrained, progress, 1.0, 1.0, **kwargs)

yolort/models/path_aggregation_network.py

@@ -57,21 +57,27 @@ def __init__(
 
         if version == "r6.0":
             init_block = SPPF(in_channels_list[2], in_channels_list[2], k=5)
+            module_version = "r4.0"
         elif version in ["r3.1", "r4.0"]:
             init_block = block(
                 in_channels_list[2], in_channels_list[2], n=depth_gain, shortcut=False
             )
+            module_version = version
         else:
             raise NotImplementedError(f"Version {version} is not implemented yet.")
 
         inner_blocks = [
             init_block,
-            Conv(in_channels_list[2], in_channels_list[1], 1, 1, version=version),
+            Conv(
+                in_channels_list[2], in_channels_list[1], 1, 1, version=module_version
+            ),
             nn.Upsample(scale_factor=2),
             block(
                 in_channels_list[2], in_channels_list[1], n=depth_gain, shortcut=False
             ),
-            Conv(in_channels_list[1], in_channels_list[0], 1, 1, version=version),
+            Conv(
+                in_channels_list[1], in_channels_list[0], 1, 1, version=module_version
+            ),
             nn.Upsample(scale_factor=2),
         ]
 
@@ -81,11 +87,15 @@ def __init__(
             block(
                 in_channels_list[1], in_channels_list[0], n=depth_gain, shortcut=False
             ),
-            Conv(in_channels_list[0], in_channels_list[0], 3, 2, version=version),
+            Conv(
+                in_channels_list[0], in_channels_list[0], 3, 2, version=module_version
+            ),
             block(
                 in_channels_list[1], in_channels_list[1], n=depth_gain, shortcut=False
             ),
-            Conv(in_channels_list[1], in_channels_list[1], 3, 2, version=version),
+            Conv(
+                in_channels_list[1], in_channels_list[1], 3, 2, version=module_version
+            ),
             block(
                 in_channels_list[2], in_channels_list[2], n=depth_gain, shortcut=False
             ),