Backbone Fix (#15)

* fixed link in docs * fixed repvgg backbone * fixed efficientnet
luxonis · Oct 9, 2024 · be74f5c · be74f5c
1 parent d78dccb
commit be74f5c
Show file tree

Hide file tree

Showing 6 changed files with 70 additions and 106 deletions.
diff --git a/luxonis_train/nodes/__init__.py b/luxonis_train/nodes/__init__.py
@@ -3,6 +3,7 @@
 from .classification_head import ClassificationHead
 from .contextspatial import ContextSpatial
 from .efficient_bbox_head import EfficientBBoxHead
+from .efficientnet import EfficientNet
 from .efficientrep import EfficientRep
 from .implicit_keypoint_bbox_head import ImplicitKeypointBBoxHead
 from .micronet import MicroNet
@@ -19,6 +20,7 @@
     "ClassificationHead",
     "ContextSpatial",
     "EfficientBBoxHead",
+    "EfficientNet",
     "EfficientRep",
     "ImplicitKeypointBBoxHead",
     "BaseNode",

diff --git a/luxonis_train/nodes/blocks/blocks.py b/luxonis_train/nodes/blocks/blocks.py
@@ -216,10 +216,7 @@ def __init__(
         kernel_size: int = 3,
         stride: int = 1,
         padding: int = 1,
-        dilation: int = 1,
         groups: int = 1,
-        padding_mode: str = "zeros",
-        deploy: bool = False,
         use_se: bool = False,
     ):
         """RepVGGBlock is a basic rep-style block, including training and deploy status
@@ -249,7 +246,6 @@ def __init__(
         """
         super().__init__()
 
-        self.deploy = deploy
         self.groups = groups
         self.in_channels = in_channels
         self.out_channels = out_channels
@@ -262,49 +258,37 @@ def __init__(
         self.nonlinearity = nn.ReLU()
 
         if use_se:
-            #   Note that RepVGG-D2se uses SE before nonlinearity. But RepVGGplus models uses SqueezeExciteBlock after nonlinearity.
+            # NOTE: that RepVGG-D2se uses SE before nonlinearity.
+            # But RepVGGplus models uses SqueezeExciteBlock after nonlinearity.
             self.se = SqueezeExciteBlock(
                 out_channels, intermediate_channels=int(out_channels // 16)
             )
         else:
-            self.se = nn.Identity()  # type: ignore
+            self.se = nn.Identity()
 
-        if deploy:
-            self.rbr_reparam = nn.Conv2d(
-                in_channels=in_channels,
-                out_channels=out_channels,
-                kernel_size=kernel_size,
-                stride=stride,
-                padding=padding,
-                dilation=dilation,
-                groups=groups,
-                bias=True,
-                padding_mode=padding_mode,
-            )
-        else:
-            self.rbr_identity = (
-                nn.BatchNorm2d(num_features=in_channels)
-                if out_channels == in_channels and stride == 1
-                else None
-            )
-            self.rbr_dense = ConvModule(
-                in_channels=in_channels,
-                out_channels=out_channels,
-                kernel_size=kernel_size,
-                stride=stride,
-                padding=padding,
-                groups=groups,
-                activation=nn.Identity(),
-            )
-            self.rbr_1x1 = ConvModule(
-                in_channels=in_channels,
-                out_channels=out_channels,
-                kernel_size=1,
-                stride=stride,
-                padding=padding_11,
-                groups=groups,
-                activation=nn.Identity(),
-            )
+        self.rbr_identity = (
+            nn.BatchNorm2d(num_features=in_channels)
+            if out_channels == in_channels and stride == 1
+            else None
+        )
+        self.rbr_dense = ConvModule(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            groups=groups,
+            activation=nn.Identity(),
+        )
+        self.rbr_1x1 = ConvModule(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            stride=stride,
+            padding=padding_11,
+            groups=groups,
+            activation=nn.Identity(),
+        )
 
     def forward(self, x: Tensor):
         if hasattr(self, "rbr_reparam"):
@@ -320,6 +304,7 @@ def forward(self, x: Tensor):
     def reparametrize(self):
         if hasattr(self, "rbr_reparam"):
             return
+
         kernel, bias = self._get_equivalent_kernel_bias()
         self.rbr_reparam = nn.Conv2d(
             in_channels=self.rbr_dense[0].in_channels,

diff --git a/luxonis_train/nodes/efficientnet.py b/luxonis_train/nodes/efficientnet.py
@@ -11,6 +11,8 @@
 
 
 class EfficientNet(BaseNode[Tensor, list[Tensor]]):
+    attach_index: int = -1
+
     def __init__(self, download_weights: bool = False, **kwargs):
         """EfficientNet backbone.
 

diff --git a/luxonis_train/nodes/efficientrep.py b/luxonis_train/nodes/efficientrep.py
@@ -17,6 +17,8 @@
 
 from .base_node import BaseNode
 
+logger = logging.getLogger(__name__)
+
 
 class EfficientRep(BaseNode[Tensor, list[Tensor]]):
     def __init__(
@@ -89,14 +91,13 @@ def __init__(
         )
 
     def set_export_mode(self, mode: bool = True) -> None:
-        """Reparametrizes instances of `RepVGGBlock` in the network.
+        """Reparametrizes instances of L{RepVGGBlock} in the network.
 
         @type mode: bool
         @param mode: Whether to set the export mode. Defaults to C{True}.
         """
         super().set_export_mode(mode)
-        logger = logging.getLogger(__name__)
-        if mode:
+        if self.export:
             logger.info("Reparametrizing EfficientRep.")
             for module in self.modules():
                 if isinstance(module, RepVGGBlock):

diff --git a/luxonis_train/nodes/mobileone.py b/luxonis_train/nodes/mobileone.py
@@ -1,38 +1,7 @@
 """MobileOne backbone.
 
-Soure: U{https://github.com/apple/ml-mobileone} @license: U{Apple
-<https://github.com/apple/ml-mobileone/blob/main/LICENSE>} @license: U{Apple
-<https://github.com/apple/ml-mobileone/blob/main/LICENSE>} @license: U{Apple
-<https://github.com/apple/ml-mobileone/blob/main/LICENSE>} @license: U{Apple
-<https://github.com/apple/ml-mobileone/blob/main/LICENSE>} @license: U{Apple
-<https://github.com/apple/ml-mobileone/blob/main/LICENSE>} @license: U{Apple
-<https://github.com/apple/ml-mobileone/blob/main/LICENSE>} @license: U{Apple
-<https://github.com/apple/ml-mobileone/blob/main/LICENSE>} @license: U{Apple
-<https://github.com/apple/ml-mobileone/blob/main/LICENSE>} @license: U{Apple
-<https://github.com/apple/ml-mobileone/blob/main/LICENSE>} @license: U{Apple
-<https://github.com/apple/ml-mobileone/blob/main/LICENSE>} @license: U{Apple
-<https://github.com/apple/ml-mobileone/blob/main/LICENSE>} @license: U{Apple
-<https://github.com/apple/ml-mobileone/blob/main/LICENSE>} @license: U{Apple
-<https://github.com/apple/ml-mobileone/blob/main/LICENSE>} @license: U{Apple
-<https://github.com/apple/ml-mobileone/blob/main/LICENSE>} @license: U{Apple
-<https://github.com/apple/ml-mobileone/blob/main/LICENSE>} @license: U{Apple
-<https://github.com/apple/ml-mobileone/blob/main/LICENSE>}
-@license: U{Apple <https://github.com/apple/ml-mobileone/blob/main/LICENSE>}
-@license: U{Apple <https://github.com/apple/ml-mobileone/blob/main/LICENSE>}
-@license: U{Apple <https://github.com/apple/ml-mobileone/blob/main/LICENSE>}
-@license: U{Apple <https://github.com/apple/ml-mobileone/blob/main/LICENSE>}
-@license: U{Apple <https://github.com/apple/ml-mobileone/blob/main/LICENSE>}
-@license: U{Apple <https://github.com/apple/ml-mobileone/blob/main/LICENSE>}
-@license: U{Apple <https://github.com/apple/ml-mobileone/blob/main/LICENSE>}
-@license: U{Apple <https://github.com/apple/ml-mobileone/blob/main/LICENSE>}
-@license: U{Apple <https://github.com/apple/ml-mobileone/blob/main/LICENSE>}
-@license: U{Apple <https://github.com/apple/ml-mobileone/blob/main/LICENSE>}
-@license: U{Apple <https://github.com/apple/ml-mobileone/blob/main/LICENSE>}
-@license: U{Apple <https://github.com/apple/ml-mobileone/blob/main/LICENSE>}
-@license: U{Apple <https://github.com/apple/ml-mobileone/blob/main/LICENSE>}
-@license: U{Apple <https://github.com/apple/ml-mobileone/blob/main/LICENSE>}
-@license: U{Apple <https://github.com/apple/ml-mobileone/blob/main/LICENSE>}
-@license: U{Apple <https://github.com/apple/ml-mobileone/blob/main/LICENSE>}
+Source: U{<https://github.com/apple/ml-mobileone>}
+@license: U{Apple<https://github.com/apple/ml-mobileone/blob/main/LICENSE>}
 """
 
 

diff --git a/luxonis_train/nodes/repvgg.py b/luxonis_train/nodes/repvgg.py
@@ -1,4 +1,5 @@
-from copy import deepcopy
+import logging
+from typing import Literal
 
 import torch.utils.checkpoint as checkpoint
 from torch import Tensor, nn
@@ -7,6 +8,8 @@
 
 from .base_node import BaseNode
 
+logger = logging.getLogger(__name__)
+
 
 class RepVGG(BaseNode):
     """Implementation of RepVGG backbone.
@@ -18,53 +21,37 @@ class RepVGG(BaseNode):
     """
 
     in_channels: int
+    attach_index: int = -1
 
     VARIANTS_SETTINGS = {
         "A0": {
             "num_blocks": [2, 4, 14, 1],
-            "num_classes": 1000,
             "width_multiplier": [0.75, 0.75, 0.75, 2.5],
         },
         "A1": {
             "num_blocks": [2, 4, 14, 1],
-            "num_classes": 1000,
             "width_multiplier": [1, 1, 1, 2.5],
         },
         "A2": {
             "num_blocks": [2, 4, 14, 1],
-            "num_classes": 1000,
             "width_multiplier": [1.5, 1.5, 1.5, 2.75],
         },
     }
 
-    def __new__(cls, **kwargs):
-        variant = kwargs.pop("variant", "A0")
-
-        if variant not in RepVGG.VARIANTS_SETTINGS.keys():
-            raise ValueError(
-                f"RepVGG model variant should be in {list(RepVGG.VARIANTS_SETTINGS.keys())}"
-            )
-
-        overrides = deepcopy(kwargs)
-        kwargs.clear()
-        kwargs.update(RepVGG.VARIANTS_SETTINGS[variant])
-        kwargs.update(overrides)
-        return cls.__new__(cls)
-
     def __init__(
         self,
-        deploy: bool = False,
+        variant: Literal["A0", "A1", "A2"] = "A0",
+        num_blocks: list[int] | None = None,
+        width_multiplier: list[float] | None = None,
         override_groups_map: dict[int, int] | None = None,
         use_se: bool = False,
         use_checkpoint: bool = False,
-        num_blocks: list[int] | None = None,
-        width_multiplier: list[float] | None = None,
         **kwargs,
     ):
         """Constructor for the RepVGG module.
 
-        @type deploy: bool
-        @param deploy: Whether to use the model in deploy mode.
+        @type variant: Literal["A0", "A1", "A2"]
+        @param variant: RepVGG model variant. Defaults to "A0".
         @type override_groups_map: dict[int, int] | None
         @param override_groups_map: Dictionary mapping layer index to number of groups.
         @type use_se: bool
@@ -77,9 +64,16 @@ def __init__(
         @param width_multiplier: Width multiplier for each stage.
         """
         super().__init__(**kwargs)
-        num_blocks = num_blocks or [2, 4, 14, 1]
-        width_multiplier = width_multiplier or [0.75, 0.75, 0.75, 2.5]
-        self.deploy = deploy
+        if variant not in self.VARIANTS_SETTINGS.keys():
+            raise ValueError(
+                f"RepVGG model variant should be one of "
+                f"{list(self.VARIANTS_SETTINGS.keys())}."
+            )
+
+        num_blocks = num_blocks or self.VARIANTS_SETTINGS[variant]["num_blocks"]
+        width_multiplier = (
+            width_multiplier or self.VARIANTS_SETTINGS[variant]["width_multiplier"]
+        )
         self.override_groups_map = override_groups_map or {}
         assert 0 not in self.override_groups_map
         self.use_se = use_se
@@ -92,7 +86,6 @@ def __init__(
             kernel_size=3,
             stride=2,
             padding=1,
-            deploy=self.deploy,
             use_se=self.use_se,
         )
         self.cur_layer_idx = 1
@@ -135,10 +128,22 @@ def _make_stage(self, planes: int, num_blocks: int, stride: int):
                     stride=stride,
                     padding=1,
                     groups=cur_groups,
-                    deploy=self.deploy,
                     use_se=self.use_se,
                 )
             )
             self.in_planes = planes
             self.cur_layer_idx += 1
         return nn.ModuleList(blocks)
+
+    def set_export_mode(self, mode: bool = True) -> None:
+        """Reparametrizes instances of L{RepVGGBlock} in the network.
+
+        @type mode: bool
+        @param mode: Whether to set the export mode. Defaults to C{True}.
+        """
+        super().set_export_mode(mode)
+        if self.export:
+            logger.info("Reparametrizing RepVGG.")
+            for module in self.modules():
+                if isinstance(module, RepVGGBlock):
+                    module.reparametrize()