Cleanup docstrings to darknet backbones (#215)

* Fixing docstrings

* Move darknetv5 to darknetv4

* Move _darknetv6 to _darknet_v6_conf

* Remove yolov5 P6 v5.0 models

* Change the base url name to match the release version

* Remove YOLOv5 P6 v5.0 in yolort.models

yolort/models/__init__.py

@@ -133,13 +133,11 @@ def yolov5s6(upstream_version: str = "r6.0", export_friendly: bool = False, **kw
     """
     Args:
         upstream_version (str): model released by the upstream YOLOv5. Possible values
-            are ["r5.0", "r6.0"]. Default: "r6.0".
+            are ["r6.0"]. Default: "r6.0".
         export_friendly (bool): Deciding whether to use (ONNX/TVM) export friendly mode.
             Default: False.
     """
-    if upstream_version == "r5.0":
-        model = YOLOv5(arch="yolov5_darknet_pan_s6_r50", **kwargs)
-    elif upstream_version == "r6.0":
+    if upstream_version == "r6.0":
         model = YOLOv5(arch="yolov5_darknet_pan_s6_r60", **kwargs)
     else:
         raise NotImplementedError("Currently only supports r5.0 and r6.0 versions")
@@ -154,13 +152,11 @@ def yolov5m6(upstream_version: str = "r6.0", export_friendly: bool = False, **kw
     """
     Args:
         upstream_version (str): model released by the upstream YOLOv5. Possible values
-            are ["r5.0", "r6.0"]. Default: "r6.0".
+            are ["r6.0"]. Default: "r6.0".
         export_friendly (bool): Deciding whether to use (ONNX/TVM) export friendly mode.
             Default: False.
     """
-    if upstream_version == "r5.0":
-        model = YOLOv5(arch="yolov5_darknet_pan_m6_r50", **kwargs)
-    elif upstream_version == "r6.0":
+    if upstream_version == "r6.0":
         model = YOLOv5(arch="yolov5_darknet_pan_m6_r60", **kwargs)
     else:
         raise NotImplementedError("Currently only supports r5.0 and r6.0 versions")

yolort/models/darknet.py

@@ -1,5 +1,5 @@
-from .darknetv5 import (
-    DarkNetV5,
+from .darknetv4 import (
+    DarkNetV4,
     darknet_s_r3_1,
     darknet_m_r3_1,
     darknet_l_r3_1,
@@ -16,7 +16,7 @@
 )
 
 __all__ = (
-    "DarkNetV5",
+    "DarkNetV4",
     "DarkNetV6",
     "darknet_s_r3_1",
     "darknet_m_r3_1",

yolort/models/darknetv5.py → yolort/models/darknetv4.py

@@ -10,7 +10,7 @@
 
 
 __all__ = [
-    "DarkNetV5",
+    "DarkNetV4",
     "darknet_s_r3_1",
     "darknet_m_r3_1",
     "darknet_l_r3_1",
@@ -29,9 +29,9 @@
 }  # TODO: add checkpoint weights
 
 
-class DarkNetV5(nn.Module):
+class DarkNetV4(nn.Module):
     """
-    DarkNetV5 main class
+    DarkNetV4 main class
 
     Args:
         depth_multiple (float): Depth multiplier
@@ -51,7 +51,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,
@@ -62,7 +62,7 @@ def __init__(
         super().__init__()
 
         assert version in ["r3.1", "r4.0"], (
-            "Currently the module version used in DarkNetV5 is r3.1 or r4.0",
+            "Currently the module version used in DarkNetV4 is r3.1 or r4.0",
         )
 
         if block is None:
@@ -135,13 +135,11 @@ def forward(self, x: Tensor) -> Tensor:
 }
 
 
-def _darknetv5(arch: str, pretrained: bool, progress: bool, *args: Any, **kwargs: Any) -> DarkNetV5:
+def _darknet_v4_conf(arch: str, pretrained: bool, progress: bool, *args: Any, **kwargs: Any) -> DarkNetV4:
     """
-    Constructs a DarkNetV5 architecture from
-    # TODO
-
+    Build a DarkNetV4 model.
     """
-    model = DarkNetV5(*args, **kwargs)
+    model = DarkNetV4(*args, **kwargs)
 
     if pretrained:
         model_url = model_urls[arch]
@@ -154,73 +152,67 @@ def _darknetv5(arch: str, pretrained: bool, progress: bool, *args: Any, **kwargs
     return model
 
 
-def darknet_s_r3_1(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> DarkNetV5:
+def darknet_s_r3_1(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> DarkNetV4:
     """
-    Constructs a DarkNetV5 with small channels, as described in release 3.1
-    # TODO
+    Constructs the DarkNet release 3.1 model with small channels.
 
     Args:
         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 _darknetv5("darknet_s_r3.1", pretrained, progress, 0.33, 0.5, "r3.1", **kwargs)
+    return _darknet_v4_conf("darknet_s_r3.1", pretrained, progress, 0.33, 0.5, version="r3.1", **kwargs)
 
 
-def darknet_m_r3_1(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> DarkNetV5:
+def darknet_m_r3_1(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> DarkNetV4:
     """
-    Constructs a DarkNetV5 with small channels, as described in release 3.1
-    # TODO
+    Constructs the DarkNet release 3.1 model with medium channels.
 
     Args:
         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 _darknetv5("darknet_m_r3.1", pretrained, progress, 0.67, 0.75, "r3.1", **kwargs)
+    return _darknet_v4_conf("darknet_m_r3.1", pretrained, progress, 0.67, 0.75, version="r3.1", **kwargs)
 
 
-def darknet_l_r3_1(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> DarkNetV5:
+def darknet_l_r3_1(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> DarkNetV4:
     """
-    Constructs a DarkNetV5 with small channels, as described in release 3.1
-    # TODO
+    Constructs the DarkNet release 3.1 model with large channels.
 
     Args:
         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 _darknetv5("darknet_l_r3.1", pretrained, progress, 1.0, 1.0, "r3.1", **kwargs)
+    return _darknet_v4_conf("darknet_l_r3.1", pretrained, progress, 1.0, 1.0, version="r3.1", **kwargs)
 
 
-def darknet_s_r4_0(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> DarkNetV5:
+def darknet_s_r4_0(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> DarkNetV4:
     """
-    Constructs a DarkNetV5 with small channels, as described in release 3.1
-    # TODO
+    Constructs the DarkNet release 4.0 model with small channels.
 
     Args:
         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 _darknetv5("darknet_s_r4.0", pretrained, progress, 0.33, 0.5, "r4.0", **kwargs)
+    return _darknet_v4_conf("darknet_s_r4.0", pretrained, progress, 0.33, 0.5, version="r4.0", **kwargs)
 
 
-def darknet_m_r4_0(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> DarkNetV5:
+def darknet_m_r4_0(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> DarkNetV4:
     """
-    Constructs a DarkNetV5 with small channels, as described in release 3.1
-    # TODO
+    Constructs the DarkNet release 4.0 model with medium channels.
 
     Args:
         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 _darknetv5("darknet_m_r4.0", pretrained, progress, 0.67, 0.75, "r4.0", **kwargs)
+    return _darknet_v4_conf("darknet_m_r4.0", pretrained, progress, 0.67, 0.75, version="r4.0", **kwargs)
 
 
-def darknet_l_r4_0(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> DarkNetV5:
+def darknet_l_r4_0(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> DarkNetV4:
     """
-    Constructs a DarkNetV5 with small channels, as described in release 3.1
-    # TODO
+    Constructs the DarkNet release 4.0 model with large channels.
 
     Args:
         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 _darknetv5("darknet_l_r4.0", pretrained, progress, 1.0, 1.0, "r4.0", **kwargs)
+    return _darknet_v4_conf("darknet_l_r4.0", pretrained, progress, 1.0, 1.0, version="r4.0", **kwargs)

yolort/models/darknetv6.py

@@ -124,11 +124,9 @@ def forward(self, x: Tensor) -> Tensor:
         return self._forward_impl(x)
 
 
-def _darknetv6(arch: str, pretrained: bool, progress: bool, *args: Any, **kwargs: Any) -> DarkNetV6:
+def _darknet_v6_conf(arch: str, pretrained: bool, progress: bool, *args: Any, **kwargs: Any) -> DarkNetV6:
     """
-    Constructs a DarkNetV6 architecture from
-    # TODO
-
+    Build a DarkNetV6 model.
     """
     model = DarkNetV6(*args, **kwargs)
 
@@ -145,47 +143,43 @@ def _darknetv6(arch: str, pretrained: bool, progress: bool, *args: Any, **kwargs
 
 def darknet_n_r6_0(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> DarkNetV6:
     """
-    Constructs a DarkNetV6 with nano channels, as described in release 6.0
-    # TODO
+    Constructs the DarkNet release 6.0 model with nano channels.
 
     Args:
         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 _darknetv6("darknet_n_r6.0", pretrained, progress, 0.33, 0.25, **kwargs)
+    return _darknet_v6_conf("darknet_n_r6.0", pretrained, progress, 0.33, 0.25, **kwargs)
 
 
 def darknet_s_r6_0(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> DarkNetV6:
     """
-    Constructs a DarkNetV6 with small channels, as described in release 6.0
-    # TODO
+    Constructs the DarkNet release 6.0 model with small channels.
 
     Args:
         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 _darknetv6("darknet_s_r6.0", pretrained, progress, 0.33, 0.5, **kwargs)
+    return _darknet_v6_conf("darknet_s_r6.0", pretrained, progress, 0.33, 0.5, **kwargs)
 
 
 def darknet_m_r6_0(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> DarkNetV6:
     """
-    Constructs a DarkNetV6 with small channels, as described in release 6.0
-    # TODO
+    Constructs the DarkNet release 6.0 model with medium channels.
 
     Args:
         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 _darknetv6("darknet_m_r6.0", pretrained, progress, 0.67, 0.75, **kwargs)
+    return _darknet_v6_conf("darknet_m_r6.0", pretrained, progress, 0.67, 0.75, **kwargs)
 
 
 def darknet_l_r6_0(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> DarkNetV6:
     """
-    Constructs a DarkNetV6 with small channels, as described in release 6.0
-    # TODO
+    Constructs the DarkNet release 6.0 model with large channels.
 
     Args:
         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 _darknetv6("darknet_l_r6.0", pretrained, progress, 1.0, 1.0, **kwargs)
+    return _darknet_v6_conf("darknet_l_r6.0", pretrained, progress, 1.0, 1.0, **kwargs)