Quality of Life improvement (#11)

.github/workflows/test.yaml

@@ -0,0 +1,37 @@
+name: Test
+on:
+  push:
+    paths-ignore:
+      - README.md
+      - .gitignore
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        pytorch-version: ["1.13", "2.0"]
+        python-version: ["3.10"]
+        include:
+          - pytorch-version: "1.13"
+            torchvision-version: "0.14"
+          - pytorch-version: "2.0"
+            torchvision-version: "0.15"
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+
+      - name: Setup Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v3
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: |
+          pip install torch==${{ matrix.pytorch-version }}.* torchvision==${{ matrix.torchvision-version }}.* --extra-index-url https://download.pytorch.org/whl/cpu
+          pip install pytest
+
+      - name: Run tests
+        run: python -m pytest -v

extras.py

@@ -7,7 +7,7 @@
 import torch
 import torch.nn.functional as F
 import torchvision.transforms.functional as TF
-from torch import nn
+from torch import Tensor, nn
 
 
 # https://github.com/pytorch/vision/blob/main/references/classification/transforms.py
@@ -19,7 +19,7 @@ def __init__(self, num_classes, p=0.5, alpha=1, inplace=False):
         self.alpha = alpha
         self.inplace = inplace
 
-    def forward(self, batch: torch.Tensor, target: torch.Tensor):
+    def forward(self, batch: Tensor, target: Tensor) -> tuple[Tensor, Tensor]:
         if not self.inplace:
             batch = batch.clone()
             target = target.clone()
@@ -35,7 +35,7 @@ def forward(self, batch: torch.Tensor, target: torch.Tensor):
         target_rolled = target.roll(1, 0)
 
         # Implemented as on mixup paper, page 3.
-        lambda_param = float(torch._sample_dirichlet(torch.tensor([self.alpha, self.alpha]))[0])
+        lambda_param = float(torch._sample_dirichlet(Tensor([self.alpha, self.alpha]))[0])
         batch_rolled.mul_(1.0 - lambda_param)
         batch.mul_(lambda_param).add_(batch_rolled)
 
@@ -53,7 +53,7 @@ def __init__(self, num_classes, p=0.5, alpha=1, inplace=False):
         self.alpha = alpha
         self.inplace = inplace
 
-    def forward(self, batch: torch.Tensor, target: torch.Tensor):
+    def forward(self, batch: Tensor, target: Tensor) -> tuple[Tensor, Tensor]:
         if not self.inplace:
             batch = batch.clone()
             target = target.clone()
@@ -129,7 +129,7 @@ def extract_backbone_weights(lightning_ckpt_path, save_name, save_dir=None):
 
 
 # Modified from YOLOv5 utils/torch_utils.py
-def profile(module: nn.Module, input: torch.Tensor = None, n: int = 10, device="cpu"):
+def profile(module: nn.Module, input: Tensor = None, n: int = 10, device="cpu"):
     from fvcore.nn import FlopCountAnalysis
 
     if input is None:

tests/test_backbones.py

@@ -1,8 +1,6 @@
-from typing import List
-
 import pytest
 import torch
-from torch import nn
+from torch import Tensor, nn
 
 from vision_toolbox import backbones
 
@@ -16,7 +14,7 @@ def inputs():
 vovnet_v2_models = [f"vovnet{x}_ese" for x in ["19_slim", 19, 39, 57, 99]]
 darknet_models = ["darknet19", "darknet53", "cspdarknet53"]
 darknet_yolov5_models = [f"darknet_yolov5{x}" for x in ("n", "s", "m", "l", "x")]
-torchvision_models = ["resnet34", "mobilenet_v2", "efficientnet_b0", "regnet_x_400mf"]
+torchvision_models = ["resnet18", "mobilenet_v2", "efficientnet_b0", "regnet_x_400mf"]
 
 all_models = vovnet_v1_models + vovnet_v2_models + darknet_models + darknet_yolov5_models + torchvision_models
 
@@ -30,7 +28,7 @@ def test_model_creation(self, name: str):
         assert isinstance(m, backbones.BaseBackbone)
 
     def test_pretrained_weights(self, name: str):
-        m = getattr(backbones, name)(pretrained=True)
+        getattr(backbones, name)(pretrained=True)
 
     def test_attributes(self, name: str):
         m = getattr(backbones, name)()
@@ -46,28 +44,32 @@ 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: torch.Tensor):
+    def test_forward(self, name: str, inputs: Tensor):
         m = getattr(backbones, name)()
         outputs = m(inputs)
 
-        assert isinstance(outputs, torch.Tensor)
+        assert isinstance(outputs, Tensor)
         assert len(outputs.shape) == 4
 
-    def test_get_feature_maps(self, name: str, inputs: torch.Tensor):
+    def test_get_feature_maps(self, name: str, inputs: Tensor):
         m = getattr(backbones, name)()
         outputs = m.get_feature_maps(inputs)
 
         assert isinstance(outputs, list)
         assert len(outputs) == len(m.out_channels_list)
         for out, out_c in zip(outputs, m.out_channels_list):
-            assert isinstance(out, torch.Tensor)
+            assert isinstance(out, Tensor)
             assert len(out.shape) == 4
             assert out.shape[1] == out_c
 
-    def test_jit_script(self, name: str):
+    def test_jit_trace(self, name: str, inputs: Tensor):
         m = getattr(backbones, name)()
-        torch.jit.script(m)
+        torch.jit.trace(m, inputs)
 
-    def test_jit_trace(self, name: str, inputs: torch.Tensor):
-        m = getattr(backbones, name)()
-        torch.jit.script(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)

vision_toolbox/backbones/base.py

@@ -1,24 +1,23 @@
 import warnings
 from abc import ABCMeta, abstractmethod
 from copy import deepcopy
-from typing import Dict, List
+from typing import Any
 
-import torch
-from torch import nn
+from torch import Tensor, nn
 from torch.hub import load_state_dict_from_url
 
 
 class BaseBackbone(nn.Module, metaclass=ABCMeta):
     # subclass only needs to implement this method
     @abstractmethod
-    def get_feature_maps(self, x: torch.Tensor) -> List[torch.Tensor]:
+    def get_feature_maps(self, x: Tensor) -> list[Tensor]:
         pass
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: Tensor) -> Tensor:
         return self.get_feature_maps(x)[-1]
 
     @classmethod
-    def from_config(cls, config: Dict, pretrained: bool = False, **kwargs):
+    def from_config(cls, config: dict[str, Any], pretrained: bool = False, **kwargs):
         config = deepcopy(config)
         weights = config.pop("weights", None)
         model = cls(**config, **kwargs)

vision_toolbox/backbones/darknet.py

@@ -3,12 +3,12 @@
 # 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, List
+from typing import Callable, Iterable
 
 import torch
-from torch import nn
+from torch import Tensor, nn
 
-from ..components import ConvBnAct
+from ..components import ConvNormAct
 from .base import BaseBackbone
 
 
@@ -30,8 +30,8 @@ class DarknetBlock(nn.Module):
     def __init__(self, in_channels: int, expansion: float = 0.5):
         super().__init__()
         mid_channels = int(in_channels * expansion)
-        self.conv1 = ConvBnAct(in_channels, mid_channels, kernel_size=1, padding=0)
-        self.conv2 = ConvBnAct(mid_channels, in_channels)
+        self.conv1 = ConvNormAct(in_channels, mid_channels, 1)
+        self.conv2 = ConvNormAct(mid_channels, in_channels)
 
     def forward(self, x):
         out = self.conv1(x)
@@ -42,23 +42,23 @@ def forward(self, x):
 class DarknetStage(nn.Sequential):
     def __init__(self, n: int, in_channels: int, out_channels: int):
         super().__init__()
-        self.conv = ConvBnAct(in_channels, out_channels, stride=2)
+        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):
         assert n > 0
         super().__init__()
-        self.conv = ConvBnAct(in_channels, out_channels, stride=2)
+        self.conv = ConvNormAct(in_channels, out_channels, stride=2)
 
         half_channels = out_channels // 2
-        self.conv1 = ConvBnAct(out_channels, half_channels, kernel_size=1, padding=0)
-        self.conv2 = ConvBnAct(out_channels, half_channels, kernel_size=1, padding=0)
+        self.conv1 = ConvNormAct(out_channels, half_channels, 1)
+        self.conv2 = ConvNormAct(out_channels, half_channels, 1)
         self.blocks = nn.Sequential(*[DarknetBlock(half_channels, expansion=1) for _ in range(n)])
-        self.out_conv = ConvBnAct(out_channels, out_channels, kernel_size=1, padding=0)
+        self.out_conv = ConvNormAct(out_channels, out_channels, 1)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: Tensor) -> Tensor:
         out = self.conv(x)
 
         # using 2 convs is faster than using 1 conv then split
@@ -74,19 +74,19 @@ class Darknet(BaseBackbone):
     def __init__(
         self,
         stem_channels: int,
-        num_blocks_list: Iterable[int],
-        num_channels_list: Iterable[int],
+        num_blocks_list: list[int],
+        num_channels_list: list[int],
         stage_cls: Callable[..., nn.Module] = DarknetStage,
     ):
         super().__init__()
         self.out_channels_list = tuple(num_channels_list)
         self.stride = 32
 
-        self.stem = ConvBnAct(3, stem_channels)
+        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):
-            self.stages.append(stage_cls(n, in_c, c) if n > 0 else ConvBnAct(in_c, c, stride=2))
+            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):
@@ -102,22 +102,22 @@ class DarknetYolov5(BaseBackbone):
     def __init__(
         self,
         stem_channels: int,
-        num_blocks_list: Iterable[int],
-        num_channels_list: Iterable[int],
+        num_blocks_list: list[int],
+        num_channels_list: list[int],
         stage_cls: Callable[..., nn.Module] = CSPDarknetStage,
     ):
         super().__init__()
         self.out_channels_list = (stem_channels,) + tuple(num_channels_list)
         self.stride = 32
 
-        self.stem = ConvBnAct(3, stem_channels, kernel_size=6, stride=2, padding=2)
+        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))
             in_c = c
 
-    def get_feature_maps(self, x: torch.Tensor) -> List[torch.Tensor]:
+    def get_feature_maps(self, x: Tensor) -> list[Tensor]:
         outputs = [self.stem(x)]
         for s in self.stages:
             outputs.append(s(outputs[-1]))
@@ -133,51 +133,51 @@ def get_feature_maps(self, x: torch.Tensor) -> List[torch.Tensor]:
     "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-da4bd7c9.pth",
+        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-7433abc1.pth",
+        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-f8225d3d.pth",
+        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-a9335553.pth",
+        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-e8dba89f.pth",
+        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-a1eb31bd.pth",
+        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-7384c740.pth",
+        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-bf388a43.pth",
+        weights="https://github.com/gau-nernst/vision-toolbox/releases/download/v0.0.1/darknet_yolov5x-0ed0c035.pth",
     ),
 }