update darknet

tests/test_backbones.py

@@ -1,8 +1,10 @@
+from functools import partial
+
 import pytest
 import torch
 from torch import Tensor, nn
 
-from vision_toolbox import backbones
+from vision_toolbox.backbones import Darknet, DarknetYOLOv5
 
 
 @pytest.fixture
@@ -19,19 +21,20 @@ def inputs():
 all_models = vovnet_v1_models + vovnet_v2_models + darknet_models + darknet_yolov5_models + torchvision_models
 
 
-@pytest.mark.parametrize("name", all_models)
-class TestBackbone:
-    def test_model_creation(self, name: str):
-        assert hasattr(backbones, name)
-        m = getattr(backbones, name)()
-        assert isinstance(m, nn.Module)
-        assert isinstance(m, backbones.BaseBackbone)
+def partial_list(fn, args_list):
+    return [partial(fn, x) for x in args_list]
+
 
-    def test_pretrained_weights(self, name: str):
-        getattr(backbones, name)(pretrained=True)
+factory_list = [
+    *partial_list(Darknet.from_config, ("darknet19", "cspdarknet53")),
+    *partial_list(DarknetYOLOv5.from_config, ("n", "l")),
+]
 
-    def test_attributes(self, name: str):
-        m = getattr(backbones, name)()
+
+@pytest.mark.parametrize("factory", factory_list)
+class TestBackbone:
+    def test_attributes(self, factory):
+        m = factory()
 
         assert hasattr(m, "out_channels_list")
         assert isinstance(m.out_channels_list, tuple)
@@ -44,15 +47,15 @@ def test_attributes(self, name: str):
         assert hasattr(m, "get_feature_maps")
         assert callable(m.get_feature_maps)
 
-    def test_forward(self, name: str, inputs: Tensor):
-        m = getattr(backbones, name)()
+    def test_forward(self, factory, inputs):
+        m = factory()
         outputs = m(inputs)
 
         assert isinstance(outputs, Tensor)
         assert len(outputs.shape) == 4
 
-    def test_get_feature_maps(self, name: str, inputs: Tensor):
-        m = getattr(backbones, name)()
+    def test_get_feature_maps(self, factory, inputs):
+        m = factory()
         outputs = m.get_feature_maps(inputs)
 
         assert isinstance(outputs, list)
@@ -62,14 +65,17 @@ def test_get_feature_maps(self, name: str, inputs: Tensor):
             assert len(out.shape) == 4
             assert out.shape[1] == out_c
 
-    def test_jit_trace(self, name: str, inputs: Tensor):
-        m = getattr(backbones, name)()
+    def test_pretrained(self, factory):
+        factory(pretrained=True)
+
+    def test_jit_trace(self, factory, inputs):
+        m = factory()
         torch.jit.trace(m, inputs)
 
 
-@pytest.mark.skipif(not hasattr(torch, "compile"), reason="torch.compile() is not available")
-@pytest.mark.parametrize("name", ["vovnet39", "vovnet19_ese", "darknet19", "cspdarknet53", "darknet_yolov5n"])
-def test_compile(name: str, inputs: Tensor):
-    m = getattr(backbones, name)()
-    m_compiled = torch.compile(m)
-    m_compiled(inputs)
+# @pytest.mark.skipif(not hasattr(torch, "compile"), reason="torch.compile() is not available")
+# @pytest.mark.parametrize("name", ["vovnet39", "vovnet19_ese", "darknet19", "cspdarknet53", "darknet_yolov5n"])
+# def test_compile(name: str, inputs: Tensor):
+#     m = getattr(backbones, name)()
+#     m_compiled = torch.compile(m)
+#     m_compiled(inputs)

vision_toolbox/backbones/__init__.py

@@ -1,5 +1,5 @@
 from .base import *
-from .darknet import *
+from .darknet import Darknet, DarknetYOLOv5
 from .patchconvnet import *
 from .torchvision_models import *
 from .vit import *

vision_toolbox/backbones/base.py

@@ -1,10 +1,10 @@
-import warnings
+from __future__ import annotations
+
 from abc import ABCMeta, abstractmethod
-from copy import deepcopy
 from typing import Any
 
+import torch
 from torch import Tensor, nn
-from torch.hub import load_state_dict_from_url
 
 
 class BaseBackbone(nn.Module, metaclass=ABCMeta):
@@ -16,16 +16,11 @@ def get_feature_maps(self, x: Tensor) -> list[Tensor]:
     def forward(self, x: Tensor) -> Tensor:
         return self.get_feature_maps(x)[-1]
 
-    @classmethod
-    def from_config(cls, config: dict[str, Any], pretrained: bool = False, **kwargs):
-        config = deepcopy(config)
-        weights = config.pop("weights", None)
-        model = cls(**config, **kwargs)
-        if pretrained:
-            if weights is not None:
-                state_dict = load_state_dict_from_url(weights)
-                model.load_state_dict(state_dict)
-            else:
-                msg = "No pre-trained weights are available. Skip loading pre-trained weights"
-                warnings.warn(msg)
-        return model
+    @staticmethod
+    @abstractmethod
+    def from_config(variant: str, pretrained: bool = False, **kwargs: Any) -> BaseBackbone:
+        pass
+
+    def _load_state_dict_from_url(self, url: str) -> None:
+        state_dict = torch.hub.load_state_dict_from_url(url)
+        self.load_state_dict(state_dict)

vision_toolbox/backbones/darknet.py

@@ -3,7 +3,9 @@
 # YOLOv3: https://arxiv.org/abs/1804.02767
 # CSPNet: https://openaccess.thecvf.com/content_CVPRW_2020/papers/w28/Wang_CSPNet_A_New_Backbone_That_Can_Enhance_Learning_Capability_of_CVPRW_2020_paper.pdf
 
-from typing import Callable, Iterable
+from __future__ import annotations
+
+from typing import Callable
 
 import torch
 from torch import Tensor, nn
@@ -12,42 +14,26 @@
 from .base import BaseBackbone
 
 
-__all__ = [
-    "Darknet",
-    "DarknetYolov5",
-    "darknet19",
-    "darknet53",
-    "cspdarknet53",
-    "darknet_yolov5n",
-    "darknet_yolov5s",
-    "darknet_yolov5m",
-    "darknet_yolov5l",
-    "darknet_yolov5x",
-]
-
-
 class DarknetBlock(nn.Module):
-    def __init__(self, in_channels: int, expansion: float = 0.5):
+    def __init__(self, in_channels: int, expansion: float = 0.5) -> None:
         super().__init__()
         mid_channels = int(in_channels * expansion)
         self.conv1 = ConvNormAct(in_channels, mid_channels, 1)
         self.conv2 = ConvNormAct(mid_channels, in_channels)
 
     def forward(self, x):
-        out = self.conv1(x)
-        out = self.conv2(out)
-        return x + out
+        return x + self.conv2(self.conv1(x))
 
 
 class DarknetStage(nn.Sequential):
-    def __init__(self, n: int, in_channels: int, out_channels: int):
+    def __init__(self, n: int, in_channels: int, out_channels: int) -> None:
         super().__init__()
         self.conv = ConvNormAct(in_channels, out_channels, stride=2)
         self.blocks = nn.Sequential(*[DarknetBlock(out_channels) for _ in range(n)])
 
 
 class CSPDarknetStage(nn.Module):
-    def __init__(self, n: int, in_channels: int, out_channels: int):
+    def __init__(self, n: int, in_channels: int, out_channels: int) -> None:
         assert n > 0
         super().__init__()
         self.conv = ConvNormAct(in_channels, out_channels, stride=2)
@@ -74,47 +60,52 @@ class Darknet(BaseBackbone):
     def __init__(
         self,
         stem_channels: int,
-        num_blocks_list: list[int],
-        num_channels_list: list[int],
+        n_blocks_list: list[int],
+        out_channels_list: list[int],
         stage_cls: Callable[..., nn.Module] = DarknetStage,
     ):
         super().__init__()
-        self.out_channels_list = tuple(num_channels_list)
+        self.out_channels_list = tuple(out_channels_list)
         self.stride = 32
 
         self.stem = ConvNormAct(3, stem_channels)
         self.stages = nn.ModuleList()
         in_c = stem_channels
-        for n, c in zip(num_blocks_list, num_channels_list):
+        for n, c in zip(n_blocks_list, out_channels_list):
             self.stages.append(stage_cls(n, in_c, c) if n > 0 else ConvNormAct(in_c, c, stride=2))
             in_c = c
 
-    def get_feature_maps(self, x):
-        outputs = []
-        out = self.stem(x)
+    def get_feature_maps(self, x: Tensor) -> list[Tensor]:
+        outputs = [self.stem(x)]
         for s in self.stages:
-            out = s(out)
-            outputs.append(out)
-        return outputs
-
-
-class DarknetYolov5(BaseBackbone):
-    def __init__(
-        self,
-        stem_channels: int,
-        num_blocks_list: list[int],
-        num_channels_list: list[int],
-        stage_cls: Callable[..., nn.Module] = CSPDarknetStage,
-    ):
+            outputs.append(s(outputs[-1]))
+        return outputs[1:]
+
+    @staticmethod
+    def from_config(variant: str, pretrained: bool = False) -> Darknet:
+        n_blocks_list, stage_cls, ckpt = dict(
+            darknet19=((0, 1, 1, 2, 2), DarknetStage, "darknet19-2cb641ca.pth"),  # YOLOv2
+            darknet53=((1, 2, 8, 8, 4), DarknetStage, "darknet53-94427f5b.pth"),  # YOLOv3
+            cspdarknet53=((1, 2, 8, 8, 4), CSPDarknetStage, "cspdarknet53-3bfa0423.pth"),  # CSPNet/YOLOv4
+        )[variant]
+        m = Darknet(32, n_blocks_list, (64, 128, 256, 512, 1024), stage_cls)
+        if pretrained:
+            base_url = "https://github.com/gau-nernst/vision-toolbox/releases/download/v0.0.1/"
+            m._load_state_dict_from_url(base_url + ckpt)
+        return m
+
+
+class DarknetYOLOv5(BaseBackbone):
+    def __init__(self, stem_channels: int, n_blocks_list: list[int], out_channels_list: list[int]) -> None:
         super().__init__()
-        self.out_channels_list = (stem_channels,) + tuple(num_channels_list)
+        self.out_channels_list = (stem_channels,) + tuple(out_channels_list)
         self.stride = 32
 
         self.stem = ConvNormAct(3, stem_channels, 6, 2)
         self.stages = nn.ModuleList()
         in_c = stem_channels
-        for n, c in zip(num_blocks_list, num_channels_list):
-            self.stages.append(stage_cls(n, in_c, c))
+        for n, c in zip(n_blocks_list, out_channels_list):
+            self.stages.append(CSPDarknetStage(n, in_c, c))
             in_c = c
 
     def get_feature_maps(self, x: Tensor) -> list[Tensor]:
@@ -123,92 +114,21 @@ def get_feature_maps(self, x: Tensor) -> list[Tensor]:
             outputs.append(s(outputs[-1]))
         return outputs
 
-
-_base = dict(stem_channels=32, num_channels_list=(64, 128, 256, 512, 1024))
-_darknet_yolov5_stem_channels = 64
-_darknet_yolov5_num_blocks_list = (3, 6, 9, 3)
-_darknet_yolov5_num_channels_list = (128, 256, 512, 1024)
-configs = {
-    # from YOLOv2
-    "darknet-19": dict(
-        **_base,
-        num_blocks_list=(0, 1, 1, 2, 2),
-        weights="https://github.com/gau-nernst/vision-toolbox/releases/download/v0.0.1/darknet19-2cb641ca.pth",
-    ),
-    # from YOLOv3
-    "darknet-53": dict(
-        **_base,
-        num_blocks_list=(1, 2, 8, 8, 4),
-        weights="https://github.com/gau-nernst/vision-toolbox/releases/download/v0.0.1/darknet53-94427f5b.pth",
-    ),
-    # from CSPNet/YOLOv4
-    "cspdarknet-53": dict(
-        **_base,
-        num_blocks_list=(1, 2, 8, 8, 4),
-        stage_cls=CSPDarknetStage,
-        weights="https://github.com/gau-nernst/vision-toolbox/releases/download/v0.0.1/cspdarknet53-3bfa0423.pth",
-    ),
-    # from YOLOv5
-    "darknet-yolov5n": dict(
-        stem_channels=int(_darknet_yolov5_stem_channels / 4),
-        num_blocks_list=tuple(int(x / 3) for x in _darknet_yolov5_num_blocks_list),
-        num_channels_list=tuple(int(x / 4) for x in _darknet_yolov5_num_channels_list),
-        weights="https://github.com/gau-nernst/vision-toolbox/releases/download/v0.0.1/darknet_yolov5n-68f182f1.pth",
-    ),
-    "darknet-yolov5s": dict(
-        stem_channels=int(_darknet_yolov5_stem_channels / 2),
-        num_blocks_list=tuple(int(x / 3) for x in _darknet_yolov5_num_blocks_list),
-        num_channels_list=tuple(int(x / 2) for x in _darknet_yolov5_num_channels_list),
-        weights="https://github.com/gau-nernst/vision-toolbox/releases/download/v0.0.1/darknet_yolov5s-175f7462.pth",
-    ),
-    "darknet-yolov5m": dict(
-        stem_channels=int(_darknet_yolov5_stem_channels * 3 / 4),
-        num_blocks_list=tuple(int(x * 2 / 3) for x in _darknet_yolov5_num_blocks_list),
-        num_channels_list=tuple(int(x * 3 / 4) for x in _darknet_yolov5_num_channels_list),
-        weights="https://github.com/gau-nernst/vision-toolbox/releases/download/v0.0.1/darknet_yolov5m-9866aa40.pth",
-    ),
-    "darknet-yolov5l": dict(
-        stem_channels=_darknet_yolov5_stem_channels,
-        num_blocks_list=_darknet_yolov5_num_blocks_list,
-        num_channels_list=_darknet_yolov5_num_channels_list,
-        weights="https://github.com/gau-nernst/vision-toolbox/releases/download/v0.0.1/darknet_yolov5l-8e25d388.pth",
-    ),
-    "darknet-yolov5x": dict(
-        stem_channels=int(_darknet_yolov5_stem_channels * 5 / 4),
-        num_blocks_list=tuple(int(x * 4 / 3) for x in _darknet_yolov5_num_blocks_list),
-        num_channels_list=tuple(int(x * 5 / 4) for x in _darknet_yolov5_num_channels_list),
-        weights="https://github.com/gau-nernst/vision-toolbox/releases/download/v0.0.1/darknet_yolov5x-0ed0c035.pth",
-    ),
-}
-
-
-def darknet19(pretrained=False, **kwargs):
-    return Darknet.from_config(configs["darknet-19"], pretrained=pretrained, **kwargs)
-
-
-def darknet53(pretrained=False, **kwargs):
-    return Darknet.from_config(configs["darknet-53"], pretrained=pretrained, **kwargs)
-
-
-def cspdarknet53(pretrained=False, **kwargs):
-    return Darknet.from_config(configs["cspdarknet-53"], pretrained=pretrained, **kwargs)
-
-
-def darknet_yolov5n(pretrained=False, **kwargs):
-    return DarknetYolov5.from_config(configs["darknet-yolov5n"], pretrained=pretrained, **kwargs)
-
-
-def darknet_yolov5s(pretrained=False, **kwargs):
-    return DarknetYolov5.from_config(configs["darknet-yolov5s"], pretrained=pretrained, **kwargs)
-
-
-def darknet_yolov5m(pretrained=False, **kwargs):
-    return DarknetYolov5.from_config(configs["darknet-yolov5m"], pretrained=pretrained, **kwargs)
-
-
-def darknet_yolov5l(pretrained=False, **kwargs):
-    return DarknetYolov5.from_config(configs["darknet-yolov5l"], pretrained=pretrained, **kwargs)
-
-
-def darknet_yolov5x(pretrained=False, **kwargs):
-    return DarknetYolov5.from_config(configs["darknet-yolov5x"], pretrained=pretrained, **kwargs)
+    @staticmethod
+    def from_config(variant: str, pretrained: bool = False) -> DarknetYOLOv5:
+        depth_scale, width_scale, ckpt = dict(
+            n=(1 / 3, 1 / 4, "darknet_yolov5n-68f182f1.pth"),
+            s=(1 / 3, 1 / 2, "darknet_yolov5s-175f7462.pth"),
+            m=(2 / 3, 3 / 4, "darknet_yolov5m-9866aa40.pth"),
+            l=(1, 1, "darknet_yolov5l-8e25d388.pth"),
+            x=(4 / 3, 5 / 4, "darknet_yolov5x-0ed0c035.pth"),
+        )[variant]
+        m = DarknetYOLOv5(
+            int(64 * width_scale),
+            [int(d * depth_scale) for d in (3, 6, 9, 3)],
+            [int(w * width_scale) for w in (128, 256, 512, 1024)],
+        )
+        if pretrained:
+            base_url = "https://github.com/gau-nernst/vision-toolbox/releases/download/v0.0.1/"
+            m._load_state_dict_from_url(base_url + ckpt)
+        return m