summaryrefslogtreecommitdiff
path: root/text_recognizer/networks/encoders
diff options
context:
space:
mode:
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(