diff --git a/luxonis_train/nodes/__init__.py b/luxonis_train/nodes/__init__.py index 954db2be..9a506c1f 100644 --- 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 index f4bd0172..4ab2ad2d 100644 --- 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 index 0b0aedde..57b52d09 100644 --- 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 index ccff4189..4e92222f 100644 --- 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 index 14e6e02b..b1658eb4 100644 --- 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 -} @license: U{Apple -} @license: U{Apple -} @license: U{Apple -} @license: U{Apple -} @license: U{Apple -} @license: U{Apple -} @license: U{Apple -} @license: U{Apple -} @license: U{Apple -} @license: U{Apple -} @license: U{Apple -} @license: U{Apple -} @license: U{Apple -} @license: U{Apple -} @license: U{Apple -} -@license: U{Apple } -@license: U{Apple } -@license: U{Apple } -@license: U{Apple } -@license: U{Apple } -@license: U{Apple } -@license: U{Apple } -@license: U{Apple } -@license: U{Apple } -@license: U{Apple } -@license: U{Apple } -@license: U{Apple } -@license: U{Apple } -@license: U{Apple } -@license: U{Apple } -@license: U{Apple } +Source: U{} +@license: U{Apple} """ diff --git a/luxonis_train/nodes/repvgg.py b/luxonis_train/nodes/repvgg.py index 44579fa5..f488a68c 100644 --- 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()