From 2951676ad13a288be0057fd2e1344395b3a61765 Mon Sep 17 00:00:00 2001 From: Jernej Sabadin Date: Tue, 19 Nov 2024 13:07:37 +0100 Subject: [PATCH] initialize_weights configurable param --- luxonis_train/nodes/README.md | 59 ++++++++++--------- .../backbones/efficientrep/efficientrep.py | 6 +- luxonis_train/nodes/blocks/blocks.py | 13 ---- .../nodes/heads/efficient_bbox_head.py | 6 +- .../nodes/necks/reppan_neck/reppan_neck.py | 6 +- 5 files changed, 46 insertions(+), 44 deletions(-) diff --git a/luxonis_train/nodes/README.md b/luxonis_train/nodes/README.md index cef35029..7c7b53c4 100644 --- a/luxonis_train/nodes/README.md +++ b/luxonis_train/nodes/README.md @@ -82,16 +82,17 @@ Adapted from [here](https://arxiv.org/pdf/2209.02976.pdf). **Parameters:** -| Key | Type | Default value | Description | -| ------------------ | ----------------------------------------------------------------- | --------------------------- | -------------------------------------------------------------------------- | -| `variant` | `Literal["n", "nano", "s", "small", "m", "medium", "l", "large"]` | `"nano"` | Variant of the network | -| `channels_list` | `list[int]` | \[64, 128, 256, 512, 1024\] | List of number of channels for each block | -| `n_repeats` | `list[int]` | \[1, 6, 12, 18, 6\] | List of number of repeats of `RepVGGBlock` | -| `depth_mul` | `float` | `0.33` | Depth multiplier | -| `width_mul` | `float` | `0.25` | Width multiplier | -| `block` | `Literal["RepBlock", "CSPStackRepBlock"]` | `"RepBlock"` | Base block used | -| `csp_e` | `float` | `0.5` | Factor for intermediate channels when block is set to `"CSPStackRepBlock"` | -| `download_weights` | `bool` | `True` | If True download weights from COCO (if available for specified variant) | +| Key | Type | Default value | Description | +| -------------------- | ----------------------------------------------------------------- | --------------------------- | -------------------------------------------------------------------------- | +| `variant` | `Literal["n", "nano", "s", "small", "m", "medium", "l", "large"]` | `"nano"` | Variant of the network | +| `channels_list` | `list[int]` | \[64, 128, 256, 512, 1024\] | List of number of channels for each block | +| `n_repeats` | `list[int]` | \[1, 6, 12, 18, 6\] | List of number of repeats of `RepVGGBlock` | +| `depth_mul` | `float` | `0.33` | Depth multiplier | +| `width_mul` | `float` | `0.25` | Width multiplier | +| `block` | `Literal["RepBlock", "CSPStackRepBlock"]` | `"RepBlock"` | Base block used | +| `csp_e` | `float` | `0.5` | Factor for intermediate channels when block is set to `"CSPStackRepBlock"` | +| `download_weights` | `bool` | `True` | If True download weights from COCO (if available for specified variant) | +| `initialize_weights` | `bool` | `True` | If True, initialize weights. | ### RexNetV1_lite @@ -175,17 +176,18 @@ Adapted from [here](https://arxiv.org/pdf/2209.02976.pdf). **Parameters:** -| Key | Type | Default value | Description | -| ------------------ | ----------------------------------------------------------------- | -------------------------------- | ------------------------------------------------------------------------------- | -| `variant` | `Literal["n", "nano", "s", "small", "m", "medium", "l", "large"]` | `"nano"` | Variant of the network | -| `n_heads` | `Literal[2,3,4]` | `3` | Number of output heads. Should be same also on the connected head in most cases | -| `channels_list` | `list[int]` | `[256, 128, 128, 256, 256, 512]` | List of number of channels for each block | -| `n_repeats` | `list[int]` | `[12, 12, 12, 12]` | List of number of repeats of `RepVGGBlock` | -| `depth_mul` | `float` | `0.33` | Depth multiplier | -| `width_mul` | `float` | `0.25` | Width multiplier | -| `block` | `Literal["RepBlock", "CSPStackRepBlock"]` | `"RepBlock"` | Base block used | -| `csp_e` | `float` | `0.5` | Factor for intermediate channels when block is set to `"CSPStackRepBlock"` | -| `download_weights` | `bool` | `False` | If True download weights from COCO (if available for specified variant) | +| Key | Type | Default value | Description | +| -------------------- | ----------------------------------------------------------------- | -------------------------------- | ------------------------------------------------------------------------------- | +| `variant` | `Literal["n", "nano", "s", "small", "m", "medium", "l", "large"]` | `"nano"` | Variant of the network | +| `n_heads` | `Literal[2,3,4]` | `3` | Number of output heads. Should be same also on the connected head in most cases | +| `channels_list` | `list[int]` | `[256, 128, 128, 256, 256, 512]` | List of number of channels for each block | +| `n_repeats` | `list[int]` | `[12, 12, 12, 12]` | List of number of repeats of `RepVGGBlock` | +| `depth_mul` | `float` | `0.33` | Depth multiplier | +| `width_mul` | `float` | `0.25` | Width multiplier | +| `block` | `Literal["RepBlock", "CSPStackRepBlock"]` | `"RepBlock"` | Base block used | +| `csp_e` | `float` | `0.5` | Factor for intermediate channels when block is set to `"CSPStackRepBlock"` | +| `download_weights` | `bool` | `False` | If True download weights from COCO (if available for specified variant) | +| `initialize_weights` | `bool` | `True` | If True, initialize weights. | ## Heads @@ -217,13 +219,14 @@ Adapted from [here](https://arxiv.org/pdf/2209.02976.pdf). **Parameters:** -| Key | Type | Default value | Description | -| ------------------ | ------- | ------------- | --------------------------------------------------------------------- | -| `n_heads` | `bool` | `3` | Number of output heads | -| `conf_thres` | `float` | `0.25` | Confidence threshold for non-maxima-suppression (used for evaluation) | -| `iou_thres` | `float` | `0.45` | `IoU` threshold for non-maxima-suppression (used for evaluation) | -| `max_det` | `int` | `300` | Maximum number of detections retained after NMS | -| `download_weights` | `bool` | `False` | If True download weights from COCO | +| Key | Type | Default value | Description | +| -------------------- | ------- | ------------- | --------------------------------------------------------------------- | +| `n_heads` | `bool` | `3` | Number of output heads | +| `conf_thres` | `float` | `0.25` | Confidence threshold for non-maxima-suppression (used for evaluation) | +| `iou_thres` | `float` | `0.45` | `IoU` threshold for non-maxima-suppression (used for evaluation) | +| `max_det` | `int` | `300` | Maximum number of detections retained after NMS | +| `download_weights` | `bool` | `False` | If True download weights from COCO | +| `initialize_weights` | `bool` | `True` | If True, initialize weights. | ### `EfficientKeypointBBoxHead` diff --git a/luxonis_train/nodes/backbones/efficientrep/efficientrep.py b/luxonis_train/nodes/backbones/efficientrep/efficientrep.py index 121ac1bc..340cc10a 100644 --- a/luxonis_train/nodes/backbones/efficientrep/efficientrep.py +++ b/luxonis_train/nodes/backbones/efficientrep/efficientrep.py @@ -30,6 +30,7 @@ def __init__( block: Literal["RepBlock", "CSPStackRepBlock"] | None = None, csp_e: float | None = None, download_weights: bool = True, + initialize_weights: bool = True, **kwargs: Any, ): """Implementation of the EfficientRep backbone. Supports the @@ -65,6 +66,8 @@ def __init__( overrides the variant value. @type download_weights: bool @param download_weights: If True download weights from COCO (if available for specified variant). Defaults to True. + @type initialize_weights: bool + @param initialize_weights: If True, initialize weights of the model. """ super().__init__(**kwargs) @@ -125,7 +128,8 @@ def __init__( ) ) - self.initialize_weights() + if initialize_weights: + self.initialize_weights() if download_weights and var.weights_path: self.load_checkpoint(var.weights_path) diff --git a/luxonis_train/nodes/blocks/blocks.py b/luxonis_train/nodes/blocks/blocks.py index 9d63853f..25bea7c5 100644 --- a/luxonis_train/nodes/blocks/blocks.py +++ b/luxonis_train/nodes/blocks/blocks.py @@ -56,19 +56,6 @@ def __init__(self, n_classes: int, in_channels: int): prior_prob = 1e-2 self._initialize_weights_and_biases(prior_prob) - self.initialize_weights() - - def initialize_weights(self): - for m in self.modules(): - if isinstance(m, nn.Conv2d): - pass - elif isinstance(m, nn.BatchNorm2d): - m.eps = 0.001 - m.momentum = 0.03 - elif isinstance( - m, (nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU) - ): - m.inplace = True def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]: out_feature = self.decoder(x) diff --git a/luxonis_train/nodes/heads/efficient_bbox_head.py b/luxonis_train/nodes/heads/efficient_bbox_head.py index 0081c6ce..531a294f 100644 --- a/luxonis_train/nodes/heads/efficient_bbox_head.py +++ b/luxonis_train/nodes/heads/efficient_bbox_head.py @@ -30,6 +30,7 @@ def __init__( iou_thres: float = 0.45, max_det: int = 300, download_weights: bool = False, + initialize_weights: bool = True, **kwargs: Any, ): """Head for object detection. @@ -51,6 +52,8 @@ def __init__( @type download_weights: bool @param download_weights: If True download weights from COCO. Defaults to False. + @type initialize_weights: bool + @param initialize_weights: If True, initialize weights. """ super().__init__(**kwargs) @@ -95,7 +98,8 @@ def __init__( f"output{i+1}_yolov6r2" for i in range(self.n_heads) ] - self.initialize_weights() + if initialize_weights: + self.initialize_weights() if download_weights: # TODO: Handle variants of head in a nicer way diff --git a/luxonis_train/nodes/necks/reppan_neck/reppan_neck.py b/luxonis_train/nodes/necks/reppan_neck/reppan_neck.py index 9d02ddcf..2c95890a 100644 --- a/luxonis_train/nodes/necks/reppan_neck/reppan_neck.py +++ b/luxonis_train/nodes/necks/reppan_neck/reppan_neck.py @@ -27,6 +27,7 @@ def __init__( block: Literal["RepBlock", "CSPStackRepBlock"] | None = None, csp_e: float | None = None, download_weights: bool = False, + initialize_weights: bool = True, **kwargs: Any, ): """Implementation of the RepPANNeck module. Supports the version @@ -65,6 +66,8 @@ def __init__( overrides the variant value. @type download_weights: bool @param download_weights: If True download weights from COCO (if available for specified variant). Defaults to False. + @type initialize_weights: bool + @param initialize_weights: If True, initialize weights of the model. """ super().__init__(**kwargs) @@ -165,7 +168,8 @@ def __init__( out_channels = channels_list_down_blocks[2 * i + 1] curr_n_repeats = n_repeats_down_blocks[i] - self.initialize_weights() + if initialize_weights: + self.initialize_weights() if download_weights and var.weights_path: self.load_checkpoint(var.weights_path)