From ca3c97c0dbfc535c128eb0f3fcea8885f250c96d Mon Sep 17 00:00:00 2001
From: deepage <604588292@qq.com>
Date: Thu, 14 Oct 2021 15:04:15 +0800
Subject: [PATCH] Add yolov5 nano models (#199)

* add yolov5_nano model

* Add yolov5n to yolort.models and hubconf.py

Co-authored-by: Zhiqiang Wang <zhiqwang@foxmail.com>
---
 hubconf.py                 |  2 +-
 yolort/models/__init__.py  | 23 ++++++++++++++++++++++-
 yolort/models/darknet.py   |  2 ++
 yolort/models/darknetv6.py | 16 ++++++++++++++++
 yolort/models/yolo.py      | 30 ++++++++++++++++++++++++++++++
 5 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/hubconf.py b/hubconf.py
index a583b7c1..85f5b77f 100644
--- a/hubconf.py
+++ b/hubconf.py
@@ -1,4 +1,4 @@
 # Optional list of dependencies required by the package
 dependencies = ["torch", "torchvision"]
 
-from yolort.models import yolov5s, yolov5m, yolov5l, yolov5ts
+from yolort.models import yolov5n, yolov5s, yolov5m, yolov5l, yolov5ts
diff --git a/yolort/models/__init__.py b/yolort/models/__init__.py
index 0bbf6b08..26021d44 100644
--- a/yolort/models/__init__.py
+++ b/yolort/models/__init__.py
@@ -8,7 +8,28 @@
 from .yolo import YOLO
 from .yolo_module import YOLOv5
 
-__all__ = ["YOLO", "YOLOv5", "yolov5s", "yolov5m", "yolov5l", "yolov5ts"]
+__all__ = ["YOLO", "YOLOv5", "yolov5n", "yolov5s", "yolov5m", "yolov5l", "yolov5ts"]
+
+
+def yolov5n(
+    upstream_version: str = "r6.0", export_friendly: bool = False, **kwargs: Any
+):
+    """
+    Args:
+        upstream_version (str): model released by the upstream YOLOv5. Possible values
+            are ["r6.0"]. Default: "r6.0".
+        export_friendly (bool): Deciding whether to use (ONNX/TVM) export friendly mode.
+            Default: False.
+    """
+    if upstream_version == "r6.0":
+        model = YOLOv5(arch="yolov5_darknet_pan_n_r60", **kwargs)
+    else:
+        raise NotImplementedError("Currently only supports r6.0 version")
+
+    if export_friendly:
+        _export_module_friendly(model)
+
+    return model
 
 
 def yolov5s(
diff --git a/yolort/models/darknet.py b/yolort/models/darknet.py
index fb93339f..bc185082 100644
--- a/yolort/models/darknet.py
+++ b/yolort/models/darknet.py
@@ -12,6 +12,7 @@
     darknet_s_r6_0,
     darknet_m_r6_0,
     darknet_l_r6_0,
+    darknet_n_r6_0,
 )
 
 __all__ = (
@@ -26,4 +27,5 @@
     "darknet_s_r6_0",
     "darknet_m_r6_0",
     "darknet_l_r6_0",
+    "darknet_n_r6_0",
 )
diff --git a/yolort/models/darknetv6.py b/yolort/models/darknetv6.py
index b4528e07..54140764 100644
--- a/yolort/models/darknetv6.py
+++ b/yolort/models/darknetv6.py
@@ -14,9 +14,11 @@
     "darknet_s_r6_0",
     "darknet_m_r6_0",
     "darknet_l_r6_0",
+    "darknet_n_r6_0",
 ]
 
 model_urls = {
+    "darknet_n_r6.0": None,
     "darknet_s_r6.0": None,
     "darknet_m_r6.0": None,
     "darknet_l_r6.0": None,
@@ -144,6 +146,20 @@ def _darknetv6(
     return model
 
 
+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
+
+    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)
+
+
 def darknet_s_r6_0(
     pretrained: bool = False, progress: bool = True, **kwargs: Any
 ) -> DarkNetV6:
diff --git a/yolort/models/yolo.py b/yolort/models/yolo.py
index 76e5f4be..312c7dc8 100644
--- a/yolort/models/yolo.py
+++ b/yolort/models/yolo.py
@@ -452,6 +452,36 @@ def yolov5_darknet_pan_l_r40(
     )
 
 
+def yolov5_darknet_pan_n_r60(
+    pretrained: bool = False,
+    progress: bool = True,
+    num_classes: int = 80,
+    **kwargs: Any,
+) -> YOLO:
+    r"""yolov5 nano release 6.0 model from
+    `"ultralytics/yolov5" <https://zenodo.org/badge/latestdoi/264818686>`_.
+    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
+    """
+    backbone_name = "darknet_n_r6_0"
+    weights_name = "yolov5_darknet_pan_n_r60_coco"
+    depth_multiple = 0.33
+    width_multiple = 0.25
+    version = "r6.0"
+    return build_model(
+        backbone_name,
+        depth_multiple,
+        width_multiple,
+        version,
+        weights_name,
+        pretrained=pretrained,
+        progress=progress,
+        num_classes=num_classes,
+        **kwargs,
+    )
+
+
 def yolov5_darknet_pan_s_r60(
     pretrained: bool = False,
     progress: bool = True,