diff options
Diffstat (limited to 'text_recognizer/networks/encoders')
| -rw-r--r-- | text_recognizer/networks/encoders/efficientnet/__init__.py (renamed from text_recognizer/networks/encoders/efficient_net/__init__.py) | 0 | ||||
| -rw-r--r-- | text_recognizer/networks/encoders/efficientnet/mbconv_block.py (renamed from text_recognizer/networks/encoders/efficient_net/block.py) | 15 |
2 files changed, 9 insertions, 6 deletions
diff --git a/text_recognizer/networks/encoders/efficient_net/__init__.py b/text_recognizer/networks/encoders/efficientnet/__init__.py index e69de29..e69de29 100644 --- a/text_recognizer/networks/encoders/efficient_net/__init__.py +++ b/text_recognizer/networks/encoders/efficientnet/__init__.py diff --git a/text_recognizer/networks/encoders/efficient_net/block.py b/text_recognizer/networks/encoders/efficientnet/mbconv_block.py index d9a0416..0384cd9 100644 --- a/text_recognizer/networks/encoders/efficient_net/block.py +++ b/text_recognizer/networks/encoders/efficientnet/mbconv_block.py @@ -27,14 +27,17 @@ class MBConvBlock(nn.Module): self.kernel_size = kernel_size self.bn_momentum = bn_momentum self.bn_eps = bn_eps + self.id_skip = id_skip self.has_se = se_ratio is not None and 0.0 < se_ratio < 1.0 - out_channels = in_channels * expand_ratio + + def _build(self, image_size: Tuple[int, int], in_channels: int, kernel_size: int, stride: int, expand_ratio: int) -> None: + inner_channels = in_channels * expand_ratio self._inverted_bottleneck = ( self._configure_inverted_bottleneck( image_size=image_size, in_channels=in_channels, - out_channels=out_channels, + out_channels=inner_channels, ) if expand_ratio != 1 else None @@ -43,8 +46,8 @@ class MBConvBlock(nn.Module): self._depthwise = self._configure_depthwise( image_size=image_size, in_channels=in_channels, - out_channels=out_channels, - groups=out_channels, + out_channels=inner_channels, + groups=inner_channels, kernel_size=kernel_size, stride=stride, ) @@ -52,14 +55,14 @@ class MBConvBlock(nn.Module): image_size = calculate_output_image_size(image_size, stride) self._squeeze_excite = ( self._configure_squeeze_excite( - in_channels=out_channels, out_channels=out_channels, se_ratio=se_ratio + in_channels=inner_channels, out_channels=inner_channels, se_ratio=se_ratio ) if self.has_se else None ) self._pointwise = self._configure_pointwise( - image_size=image_size, in_channels=out_channels, out_channels=out_channels + image_size=image_size, in_channels=inner_channels, out_channels=out_channels ) def _configure_inverted_bottleneck( |