diff options
Diffstat (limited to 'text_recognizer/networks/efficientnet')
| -rw-r--r-- | text_recognizer/networks/efficientnet/efficientnet.py | 47 | ||||
| -rw-r--r-- | text_recognizer/networks/efficientnet/mbconv.py | 121 |
2 files changed, 99 insertions, 69 deletions
diff --git a/text_recognizer/networks/efficientnet/efficientnet.py b/text_recognizer/networks/efficientnet/efficientnet.py index cf64bcf..2260ee2 100644 --- a/text_recognizer/networks/efficientnet/efficientnet.py +++ b/text_recognizer/networks/efficientnet/efficientnet.py @@ -1,7 +1,6 @@ """Efficientnet backbone.""" from typing import Tuple -from attrs import define, field from torch import nn, Tensor from text_recognizer.networks.efficientnet.mbconv import MBConvBlock @@ -12,13 +11,9 @@ from text_recognizer.networks.efficientnet.utils import ( ) -@define(eq=False) class EfficientNet(nn.Module): """Efficientnet without classification head.""" - def __attrs_pre_init__(self) -> None: - super().__init__() - archs = { # width, depth, dropout "b0": (1.0, 1.0, 0.2), @@ -33,32 +28,32 @@ class EfficientNet(nn.Module): "l2": (4.3, 5.3, 0.5), } - arch: str = field() - params: Tuple[float, float, float] = field(default=None, init=False) - stochastic_dropout_rate: float = field(default=0.2) - bn_momentum: float = field(default=0.99) - bn_eps: float = field(default=1.0e-3) - depth: int = field(default=7) - out_channels: int = field(default=None, init=False) - _conv_stem: nn.Sequential = field(default=None, init=False) - _blocks: nn.ModuleList = field(default=None, init=False) - _conv_head: nn.Sequential = field(default=None, init=False) - - def __attrs_post_init__(self) -> None: - """Post init configuration.""" + def __init__( + self, + arch: str, + params: Tuple[float, float, float], + stochastic_dropout_rate: float = 0.2, + bn_momentum: float = 0.99, + bn_eps: float = 1.0e-3, + depth: int = 7, + ) -> None: + super().__init__() + self.params = self._get_arch_params(arch) + self.stochastic_dropout_rate = stochastic_dropout_rate + self.bn_momentum = bn_momentum + self.bn_eps = bn_eps + self.depth = depth + self.out_channels: int + self._conv_stem: nn.Sequential + self._blocks: nn.ModuleList + self._conv_head: nn.Sequential self._build() - @depth.validator - def _check_depth(self, attribute, value: str) -> None: - if not 5 <= value <= 7: - raise ValueError(f"Depth has to be between 5 and 7, was: {value}") - - @arch.validator - def _check_arch(self, attribute, value: str) -> None: + def _get_arch_params(self, value: str) -> Tuple[float, float, float]: """Validates the efficientnet architecure.""" if value not in self.archs: raise ValueError(f"{value} not a valid architecure.") - self.params = self.archs[value] + return self.archs[value] def _build(self) -> None: """Builds the efficientnet backbone.""" diff --git a/text_recognizer/networks/efficientnet/mbconv.py b/text_recognizer/networks/efficientnet/mbconv.py index 98e9353..64debd9 100644 --- a/text_recognizer/networks/efficientnet/mbconv.py +++ b/text_recognizer/networks/efficientnet/mbconv.py @@ -1,7 +1,6 @@ """Mobile inverted residual block.""" from typing import Optional, Tuple, Union -from attrs import define, field import torch from torch import nn, Tensor import torch.nn.functional as F @@ -14,18 +13,15 @@ def _convert_stride(stride: Union[Tuple[int, int], int]) -> Tuple[int, int]: return (stride,) * 2 if isinstance(stride, int) else stride -@define(eq=False) class BaseModule(nn.Module): """Base sub module class.""" - bn_momentum: float = field() - bn_eps: float = field() - block: nn.Sequential = field(init=False) - - def __attrs_pre_init__(self) -> None: + def __init__(self, bn_momentum: float, bn_eps: float, block: nn.Sequential) -> None: super().__init__() - def __attrs_post_init__(self) -> None: + self.bn_momentum = bn_momentum + self.bn_eps = bn_eps + self.block = block self._build() def _build(self) -> None: @@ -36,12 +32,20 @@ class BaseModule(nn.Module): return self.block(x) -@define(auto_attribs=True, eq=False) class InvertedBottleneck(BaseModule): """Inverted bottleneck module.""" - in_channels: int = field() - out_channels: int = field() + def __init__( + self, + bn_momentum: float, + bn_eps: float, + block: nn.Sequential, + in_channels: int, + out_channels: int, + ) -> None: + super().__init__(bn_momentum, bn_eps, block) + self.in_channels = in_channels + self.out_channels = out_channels def _build(self) -> None: self.block = nn.Sequential( @@ -60,13 +64,22 @@ class InvertedBottleneck(BaseModule): ) -@define(auto_attribs=True, eq=False) class Depthwise(BaseModule): """Depthwise convolution module.""" - channels: int = field() - kernel_size: int = field() - stride: int = field() + def __init__( + self, + bn_momentum: float, + bn_eps: float, + block: nn.Sequential, + channels: int, + kernel_size: int, + stride: int, + ) -> None: + super().__init__(bn_momentum, bn_eps, block) + self.channels = channels + self.kernel_size = kernel_size + self.stride = stride def _build(self) -> None: self.block = nn.Sequential( @@ -85,13 +98,23 @@ class Depthwise(BaseModule): ) -@define(auto_attribs=True, eq=False) class SqueezeAndExcite(BaseModule): """Sequeeze and excite module.""" - in_channels: int = field() - channels: int = field() - se_ratio: float = field() + def __init__( + self, + bn_momentum: float, + bn_eps: float, + block: nn.Sequential, + in_channels: int, + channels: int, + se_ratio: float, + ) -> None: + super().__init__(bn_momentum, bn_eps, block) + + self.in_channels = in_channels + self.channels = channels + self.se_ratio = se_ratio def _build(self) -> None: num_squeezed_channels = max(1, int(self.in_channels * self.se_ratio)) @@ -110,12 +133,20 @@ class SqueezeAndExcite(BaseModule): ) -@define(auto_attribs=True, eq=False) class Pointwise(BaseModule): """Pointwise module.""" - in_channels: int = field() - out_channels: int = field() + def __init__( + self, + bn_momentum: float, + bn_eps: float, + block: nn.Sequential, + in_channels: int, + out_channels: int, + ) -> None: + super().__init__(bn_momentum, bn_eps, block) + self.in_channels = in_channels + self.out_channels = out_channels def _build(self) -> None: self.block = nn.Sequential( @@ -133,28 +164,36 @@ class Pointwise(BaseModule): ) -@define(eq=False) class MBConvBlock(nn.Module): """Mobile Inverted Residual Bottleneck block.""" - def __attrs_pre_init__(self) -> None: + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: Tuple[int, int], + stride: Tuple[int, int], + bn_momentum: float, + bn_eps: float, + se_ratio: float, + expand_ratio: int, + ) -> None: super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.kernel_size = kernel_size + self.stride = stride + self.bn_momentum = bn_momentum + self.bn_eps = bn_eps + self.se_ratio = se_ratio + self.expand_ratio = expand_ratio + self.pad = self._configure_padding() + self._inverted_bottleneck: Optional[InvertedBottleneck] + self._depthwise: nn.Sequential + self._squeeze_excite: nn.Sequential + self._pointwise: nn.Sequential + self._build() - in_channels: int = field() - out_channels: int = field() - kernel_size: Tuple[int, int] = field() - stride: Tuple[int, int] = field(converter=_convert_stride) - bn_momentum: float = field() - bn_eps: float = field() - se_ratio: float = field() - expand_ratio: int = field() - pad: Tuple[int, int, int, int] = field(init=False) - _inverted_bottleneck: Optional[InvertedBottleneck] = field(init=False) - _depthwise: nn.Sequential = field(init=False) - _squeeze_excite: nn.Sequential = field(init=False) - _pointwise: nn.Sequential = field(init=False) - - @pad.default def _configure_padding(self) -> Tuple[int, int, int, int]: """Set padding for convolutional layers.""" if self.stride == (2, 2): @@ -164,10 +203,6 @@ class MBConvBlock(nn.Module): ) * 2 return ((self.kernel_size - 1) // 2,) * 4 - def __attrs_post_init__(self) -> None: - """Post init configuration.""" - self._build() - def _build(self) -> None: has_se = self.se_ratio is not None and 0.0 < self.se_ratio < 1.0 inner_channels = self.in_channels * self.expand_ratio |