From 050e1bd284a173d2586ad4607e95d114691db563 Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Mon, 22 Nov 2021 22:38:43 +0100
Subject: Move efficientnet from encoder dir
---
.../networks/encoders/efficientnet/__init__.py | 2 -
.../networks/encoders/efficientnet/efficientnet.py | 118 ----------
.../networks/encoders/efficientnet/mbconv.py | 240 ---------------------
.../networks/encoders/efficientnet/utils.py | 85 --------
4 files changed, 445 deletions(-)
delete mode 100644 text_recognizer/networks/encoders/efficientnet/__init__.py
delete mode 100644 text_recognizer/networks/encoders/efficientnet/efficientnet.py
delete mode 100644 text_recognizer/networks/encoders/efficientnet/mbconv.py
delete mode 100644 text_recognizer/networks/encoders/efficientnet/utils.py
(limited to 'text_recognizer/networks/encoders/efficientnet')
diff --git a/text_recognizer/networks/encoders/efficientnet/__init__.py b/text_recognizer/networks/encoders/efficientnet/__init__.py
deleted file mode 100644
index 344233f..0000000
--- a/text_recognizer/networks/encoders/efficientnet/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-"""Efficient net."""
-from .efficientnet import EfficientNet
diff --git a/text_recognizer/networks/encoders/efficientnet/efficientnet.py b/text_recognizer/networks/encoders/efficientnet/efficientnet.py
deleted file mode 100644
index 9454514..0000000
--- a/text_recognizer/networks/encoders/efficientnet/efficientnet.py
+++ /dev/null
@@ -1,118 +0,0 @@
-"""Efficientnet backbone."""
-from typing import Tuple
-
-import attr
-from torch import nn, Tensor
-
-from .mbconv import MBConvBlock
-from .utils import (
- block_args,
- round_filters,
- round_repeats,
-)
-
-
-@attr.s(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),
- "b1": (1.0, 1.1, 0.2),
- "b2": (1.1, 1.2, 0.3),
- "b3": (1.2, 1.4, 0.3),
- "b4": (1.4, 1.8, 0.4),
- "b5": (1.6, 2.2, 0.4),
- "b6": (1.8, 2.6, 0.5),
- "b7": (2.0, 3.1, 0.5),
- "b8": (2.2, 3.6, 0.5),
- "l2": (4.3, 5.3, 0.5),
- }
-
- arch: str = attr.ib()
- params: Tuple[float, float, float] = attr.ib(default=None, init=False)
- stochastic_dropout_rate: float = attr.ib(default=0.2)
- bn_momentum: float = attr.ib(default=0.99)
- bn_eps: float = attr.ib(default=1.0e-3)
- out_channels: int = attr.ib(default=None, init=False)
- _conv_stem: nn.Sequential = attr.ib(default=None, init=False)
- _blocks: nn.ModuleList = attr.ib(default=None, init=False)
- _conv_head: nn.Sequential = attr.ib(default=None, init=False)
-
- def __attrs_post_init__(self) -> None:
- """Post init configuration."""
- self._build()
-
- @arch.validator
- def check_arch(self, attribute: attr._make.Attribute, value: str) -> None:
- """Validates the efficientnet architecure."""
- if value not in self.archs:
- raise ValueError(f"{value} not a valid architecure.")
- self.params = self.archs[value]
-
- def _build(self) -> None:
- """Builds the efficientnet backbone."""
- _block_args = block_args()
- in_channels = 1 # BW
- out_channels = round_filters(32, self.params)
- self._conv_stem = nn.Sequential(
- nn.ZeroPad2d((0, 1, 0, 1)),
- nn.Conv2d(
- in_channels=in_channels,
- out_channels=out_channels,
- kernel_size=3,
- stride=(2, 2),
- bias=False,
- ),
- nn.BatchNorm2d(
- num_features=out_channels, momentum=self.bn_momentum, eps=self.bn_eps
- ),
- nn.Mish(inplace=True),
- )
- self._blocks = nn.ModuleList([])
- for args in _block_args:
- args.in_channels = round_filters(args.in_channels, self.params)
- args.out_channels = round_filters(args.out_channels, self.params)
- num_repeats = round_repeats(args.num_repeats, self.params)
- del args.num_repeats
- for _ in range(num_repeats):
- self._blocks.append(
- MBConvBlock(
- **args, bn_momentum=self.bn_momentum, bn_eps=self.bn_eps,
- )
- )
- args.in_channels = args.out_channels
- args.stride = 1
-
- in_channels = round_filters(320, self.params)
- self.out_channels = round_filters(1280, self.params)
- self._conv_head = nn.Sequential(
- nn.Conv2d(
- in_channels, self.out_channels, kernel_size=1, stride=1, bias=False
- ),
- nn.BatchNorm2d(
- num_features=self.out_channels,
- momentum=self.bn_momentum,
- eps=self.bn_eps,
- ),
- nn.Dropout(p=self.params[-1]),
- )
-
- def extract_features(self, x: Tensor) -> Tensor:
- """Extracts the final feature map layer."""
- x = self._conv_stem(x)
- for i, block in enumerate(self._blocks):
- stochastic_dropout_rate = self.stochastic_dropout_rate
- if self.stochastic_dropout_rate:
- stochastic_dropout_rate *= i / len(self._blocks)
- x = block(x, stochastic_dropout_rate=stochastic_dropout_rate)
- x = self._conv_head(x)
- return x
-
- def forward(self, x: Tensor) -> Tensor:
- """Returns efficientnet image features."""
- return self.extract_features(x)
diff --git a/text_recognizer/networks/encoders/efficientnet/mbconv.py b/text_recognizer/networks/encoders/efficientnet/mbconv.py
deleted file mode 100644
index 4b051eb..0000000
--- a/text_recognizer/networks/encoders/efficientnet/mbconv.py
+++ /dev/null
@@ -1,240 +0,0 @@
-"""Mobile inverted residual block."""
-from typing import Optional, Tuple, Union
-
-import attr
-import torch
-from torch import nn, Tensor
-import torch.nn.functional as F
-
-from text_recognizer.networks.encoders.efficientnet.utils import stochastic_depth
-
-
-def _convert_stride(stride: Union[Tuple[int, int], int]) -> Tuple[int, int]:
- """Converts int to tuple."""
- return (stride,) * 2 if isinstance(stride, int) else stride
-
-
-@attr.s(eq=False)
-class BaseModule(nn.Module):
- """Base sub module class."""
-
- bn_momentum: float = attr.ib()
- bn_eps: float = attr.ib()
- block: nn.Sequential = attr.ib(init=False)
-
- def __attrs_pre_init__(self) -> None:
- super().__init__()
-
- def __attrs_post_init__(self) -> None:
- self._build()
-
- def _build(self) -> None:
- pass
-
- def forward(self, x: Tensor) -> Tensor:
- """Forward pass."""
- return self.block(x)
-
-
-@attr.s(auto_attribs=True, eq=False)
-class InvertedBottleneck(BaseModule):
- """Inverted bottleneck module."""
-
- in_channels: int = attr.ib()
- out_channels: int = attr.ib()
-
- def _build(self) -> None:
- self.block = nn.Sequential(
- nn.Conv2d(
- in_channels=self.in_channels,
- out_channels=self.out_channels,
- kernel_size=1,
- bias=False,
- ),
- nn.BatchNorm2d(
- num_features=self.out_channels,
- momentum=self.bn_momentum,
- eps=self.bn_eps,
- ),
- nn.Mish(inplace=True),
- )
-
-
-@attr.s(auto_attribs=True, eq=False)
-class Depthwise(BaseModule):
- """Depthwise convolution module."""
-
- channels: int = attr.ib()
- kernel_size: int = attr.ib()
- stride: int = attr.ib()
-
- def _build(self) -> None:
- self.block = nn.Sequential(
- nn.Conv2d(
- in_channels=self.channels,
- out_channels=self.channels,
- kernel_size=self.kernel_size,
- stride=self.stride,
- groups=self.channels,
- bias=False,
- ),
- nn.BatchNorm2d(
- num_features=self.channels, momentum=self.bn_momentum, eps=self.bn_eps
- ),
- nn.Mish(inplace=True),
- )
-
-
-@attr.s(auto_attribs=True, eq=False)
-class SqueezeAndExcite(BaseModule):
- """Sequeeze and excite module."""
-
- in_channels: int = attr.ib()
- channels: int = attr.ib()
- se_ratio: float = attr.ib()
-
- def _build(self) -> None:
- num_squeezed_channels = max(1, int(self.in_channels * self.se_ratio))
- self.block = nn.Sequential(
- nn.Conv2d(
- in_channels=self.channels,
- out_channels=num_squeezed_channels,
- kernel_size=1,
- ),
- nn.Mish(inplace=True),
- nn.Conv2d(
- in_channels=num_squeezed_channels,
- out_channels=self.channels,
- kernel_size=1,
- ),
- )
-
-
-@attr.s(auto_attribs=True, eq=False)
-class Pointwise(BaseModule):
- """Pointwise module."""
-
- in_channels: int = attr.ib()
- out_channels: int = attr.ib()
-
- def _build(self) -> None:
- self.block = nn.Sequential(
- nn.Conv2d(
- in_channels=self.in_channels,
- out_channels=self.out_channels,
- kernel_size=1,
- bias=False,
- ),
- nn.BatchNorm2d(
- num_features=self.out_channels,
- momentum=self.bn_momentum,
- eps=self.bn_eps,
- ),
- )
-
-
-@attr.s(eq=False)
-class MBConvBlock(nn.Module):
- """Mobile Inverted Residual Bottleneck block."""
-
- def __attrs_pre_init__(self) -> None:
- super().__init__()
-
- in_channels: int = attr.ib()
- out_channels: int = attr.ib()
- kernel_size: Tuple[int, int] = attr.ib()
- stride: Tuple[int, int] = attr.ib(converter=_convert_stride)
- bn_momentum: float = attr.ib()
- bn_eps: float = attr.ib()
- se_ratio: float = attr.ib()
- expand_ratio: int = attr.ib()
- pad: Tuple[int, int, int, int] = attr.ib(init=False)
- _inverted_bottleneck: Optional[InvertedBottleneck] = attr.ib(init=False)
- _depthwise: nn.Sequential = attr.ib(init=False)
- _squeeze_excite: nn.Sequential = attr.ib(init=False)
- _pointwise: nn.Sequential = attr.ib(init=False)
-
- @pad.default
- def _configure_padding(self) -> Tuple[int, int, int, int]:
- """Set padding for convolutional layers."""
- if self.stride == (2, 2):
- return ((self.kernel_size - 1) // 2 - 1, (self.kernel_size - 1) // 2,) * 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
- self._inverted_bottleneck = (
- InvertedBottleneck(
- in_channels=self.in_channels,
- out_channels=inner_channels,
- bn_momentum=self.bn_momentum,
- bn_eps=self.bn_eps,
- )
- if self.expand_ratio != 1
- else None
- )
-
- self._depthwise = Depthwise(
- channels=inner_channels,
- kernel_size=self.kernel_size,
- stride=self.stride,
- bn_momentum=self.bn_momentum,
- bn_eps=self.bn_eps,
- )
-
- self._squeeze_excite = (
- SqueezeAndExcite(
- in_channels=self.in_channels,
- channels=inner_channels,
- se_ratio=self.se_ratio,
- bn_momentum=self.bn_momentum,
- bn_eps=self.bn_eps,
- )
- if has_se
- else None
- )
-
- self._pointwise = Pointwise(
- in_channels=inner_channels,
- out_channels=self.out_channels,
- bn_momentum=self.bn_momentum,
- bn_eps=self.bn_eps,
- )
-
- def _stochastic_depth(
- self, x: Tensor, residual: Tensor, stochastic_dropout_rate: Optional[float]
- ) -> Tensor:
- if self.stride == (1, 1) and self.in_channels == self.out_channels:
- if stochastic_dropout_rate:
- x = stochastic_depth(
- x, p=stochastic_dropout_rate, training=self.training
- )
- x += residual
- return x
-
- def forward(
- self, x: Tensor, stochastic_dropout_rate: Optional[float] = None
- ) -> Tensor:
- """Forward pass."""
- residual = x
- if self._inverted_bottleneck is not None:
- x = self._inverted_bottleneck(x)
-
- x = F.pad(x, self.pad)
- x = self._depthwise(x)
-
- if self._squeeze_excite is not None:
- x_squeezed = F.adaptive_avg_pool2d(x, 1)
- x_squeezed = self._squeeze_excite(x)
- x = torch.tanh(F.softplus(x_squeezed)) * x
-
- x = self._pointwise(x)
-
- # Stochastic depth
- x = self._stochastic_depth(x, residual, stochastic_dropout_rate)
- return x
diff --git a/text_recognizer/networks/encoders/efficientnet/utils.py b/text_recognizer/networks/encoders/efficientnet/utils.py
deleted file mode 100644
index 2b1aebb..0000000
--- a/text_recognizer/networks/encoders/efficientnet/utils.py
+++ /dev/null
@@ -1,85 +0,0 @@
-"""Util functions for efficient net."""
-import math
-from typing import List, Tuple
-
-from omegaconf import DictConfig, OmegaConf
-import torch
-from torch import Tensor
-
-
-def stochastic_depth(x: Tensor, p: float, training: bool) -> Tensor:
- """Stochastic connection.
-
- Drops the entire convolution with a given survival probability.
-
- Args:
- x (Tensor): Input tensor.
- p (float): Survival probability between 0.0 and 1.0.
- training (bool): The running mode.
-
- Shapes:
- - x: :math: `(B, C, W, H)`.
- - out: :math: `(B, C, W, H)`.
-
- where B is the batch size, C is the number of channels, W is the width, and H
- is the height.
-
- Returns:
- out (Tensor): Output after drop connection.
- """
- assert 0.0 <= p <= 1.0, "p must be in range of [0, 1]"
-
- if not training:
- return x
-
- bsz = x.shape[0]
- survival_prob = 1 - p
-
- # Generate a binary tensor mask according to probability (p for 0, 1-p for 1)
- random_tensor = survival_prob
- random_tensor += torch.rand([bsz, 1, 1, 1]).type_as(x)
- binary_tensor = torch.floor(random_tensor)
-
- out = x / survival_prob * binary_tensor
- return out
-
-
-def round_filters(filters: int, arch: Tuple[float, float, float]) -> int:
- """Returns the number output filters for a block."""
- multiplier = arch[0]
- divisor = 8
- filters *= multiplier
- new_filters = max(divisor, (filters + divisor // 2) // divisor * divisor)
- if new_filters < 0.9 * filters:
- new_filters += divisor
- return int(new_filters)
-
-
-def round_repeats(repeats: int, arch: Tuple[float, float, float]) -> int:
- """Returns how many times a layer should be repeated in a block."""
- return int(math.ceil(arch[1] * repeats))
-
-
-def block_args() -> List[DictConfig]:
- """Returns arguments for each efficientnet block."""
- keys = [
- "num_repeats",
- "kernel_size",
- "stride",
- "expand_ratio",
- "in_channels",
- "out_channels",
- "se_ratio",
- ]
- args = [
- [1, 3, (1, 1), 1, 32, 16, 0.25],
- [2, 3, (2, 2), 6, 16, 24, 0.25],
- [2, 5, (2, 2), 6, 24, 40, 0.25],
- [3, 3, (2, 2), 6, 40, 80, 0.25],
- [3, 5, (1, 1), 6, 80, 112, 0.25],
- [1, 3, (1, 1), 6, 112, 320, 0.25],
- ]
- block_args_ = []
- for row in args:
- block_args_.append(OmegaConf.create(dict(zip(keys, row))))
- return block_args_
--
cgit v1.2.3-70-g09d2