diff options
Diffstat (limited to 'text_recognizer/networks')
4 files changed, 0 insertions, 250 deletions
diff --git a/text_recognizer/networks/transformer/__init__.py b/text_recognizer/networks/transformer/__init__.py index d9e63ef..a3f3011 100644 --- a/text_recognizer/networks/transformer/__init__.py +++ b/text_recognizer/networks/transformer/__init__.py @@ -1,3 +1 @@ """Transformer modules.""" -from .nystromer.nystromer import Nystromer -from .vit import ViT diff --git a/text_recognizer/networks/transformer/nystromer/__init__.py b/text_recognizer/networks/transformer/nystromer/__init__.py deleted file mode 100644 index ea2c6fc..0000000 --- a/text_recognizer/networks/transformer/nystromer/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -"""Nyströmer module.""" -from .nystromer import Nystromer diff --git a/text_recognizer/networks/transformer/nystromer/attention.py b/text_recognizer/networks/transformer/nystromer/attention.py deleted file mode 100644 index 695a0d7..0000000 --- a/text_recognizer/networks/transformer/nystromer/attention.py +++ /dev/null @@ -1,182 +0,0 @@ -"""Nyströmer encoder. - -Efficient attention module that reduces the complexity of the attention module from -O(n**2) to O(n). The attention matrix is assumed low rank and thus the information -can be represented by a smaller matrix. - -Stolen from: - https://github.com/lucidrains/nystrom-attention/blob/main/nystrom_attention/nystrom_attention.py - -""" -from math import ceil -from typing import Optional, Tuple, Union - -from einops import rearrange, reduce -import torch -from torch import einsum, nn, Tensor -from torch.nn import functional as F - - -def moore_penrose_inverse(x: Tensor, iters: int = 6) -> Tensor: - """Moore-Penrose pseudoinverse.""" - x_abs = torch.abs(x) - col = x_abs.sum(dim=-1) - row = x_abs.sum(dim=-2) - z = rearrange(x, "... i j -> ... j i") / (torch.max(col) * torch.max(row)) - - I = torch.eye(x.shape[-1], device=x.device) - I = rearrange(I, "i j -> () i j") - - for _ in range(iters): - xz = x @ z - z = 0.25 * z @ (13 * I - (xz @ (15 * I - (xz @ (7 * I - xz))))) - return z - - -class NystromAttention(nn.Module): - def __init__( - self, - dim: int, - dim_head: int = 64, - num_heads: int = 8, - num_landmarks: int = 256, - inverse_iter: int = 6, - residual: bool = True, - residual_conv_kernel: int = 13, - eps: float = 1.0e-8, - dropout_rate: float = 0.0, - ): - super().__init__() - self.dim = dim - self.residual = None - self.eps = eps - self.num_heads = num_heads - inner_dim = self.num_heads * dim_head - self.num_landmarks = num_landmarks - self.inverse_iter = inverse_iter - self.scale = dim_head ** -0.5 - - self.qkv_fn = nn.Linear(dim, 3 * inner_dim, bias=False) - self.fc_out = nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout_rate)) - - if residual: - self.residual = nn.Conv2d( - in_channels=num_heads, - out_channels=num_heads, - kernel_size=(residual_conv_kernel, 1), - padding=(residual_conv_kernel // 2, 0), - groups=num_heads, - bias=False, - ) - - @staticmethod - def _pad_sequence( - x: Tensor, mask: Optional[Tensor], n: int, m: int - ) -> Tuple[Tensor, Tensor]: - """Pad sequence.""" - padding = m - (n % m) - x = F.pad(x, (0, 0, padding, 0), value=0) - mask = F.pad(mask, (padding, 0), value=False) if mask is not None else mask - return x, mask - - def _compute_landmarks( - self, q: Tensor, k: Tensor, mask: Optional[Tensor], n: int, m: int - ) -> Tuple[Tensor, Tensor, Optional[Tensor]]: - """Compute landmarks of the attention matrix.""" - divisor = ceil(n / m) - landmark_einops_eq = "... (n l) d -> ... n d" - q_landmarks = reduce(q, landmark_einops_eq, "sum", l=divisor) - k_landmarks = reduce(k, landmark_einops_eq, "sum", l=divisor) - - mask_landmarks = None - if mask is not None: - mask_landmarks_sum = reduce(mask, "... (n l) -> ... n", "sum", l=divisor) - divisor = mask_landmarks_sum[..., None] + self.eps - mask_landmarks = mask_landmarks_sum > 0 - - q_landmarks /= divisor - k_landmarks /= divisor - - return q_landmarks, k_landmarks, mask_landmarks - - @staticmethod - def _compute_similarities( - q: Tensor, - k: Tensor, - q_landmarks: Tensor, - k_landmarks: Tensor, - mask: Optional[Tensor], - mask_landmarks: Optional[Tensor], - ) -> Tuple[Tensor, Tensor, Tensor]: - einops_eq = "... i d, ... j d -> ... i j" - sim1 = einsum(einops_eq, q, k_landmarks) - sim2 = einsum(einops_eq, q_landmarks, k_landmarks) - sim3 = einsum(einops_eq, q_landmarks, k) - - if mask is not None and mask_landmarks is not None: - mask_value = -torch.finfo(q.type).max - sim1.masked_fill_( - ~(mask[..., None] * mask_landmarks[..., None, :]), mask_value - ) - sim2.masked_fill_( - ~(mask_landmarks[..., None] * mask_landmarks[..., None, :]), mask_value - ) - sim3.masked_fill_( - ~(mask_landmarks[..., None] * mask[..., None, :]), mask_value - ) - - return sim1, sim2, sim3 - - def _nystrom_attention( - self, - q: Tensor, - k: Tensor, - v: Tensor, - mask: Optional[Tensor], - n: int, - m: int, - return_attn: bool, - ) -> Tuple[Tensor, Optional[Tensor]]: - q_landmarks, k_landmarks, mask_landmarks = self._compute_landmarks( - q, k, mask, n, m - ) - sim1, sim2, sim3 = self._compute_similarities( - q, k, q_landmarks, k_landmarks, mask, mask_landmarks - ) - - # Compute attention - attn1, attn2, attn3 = map(lambda t: t.softmax(dim=-1), (sim1, sim2, sim3)) - attn2_inv = moore_penrose_inverse(attn2, self.inverse_iter) - - out = (attn1 @ attn2_inv) @ (attn3 @ v) - - if return_attn: - return out, attn1 @ attn2_inv @ attn3 - return out, None - - def forward( - self, x: Tensor, mask: Optional[Tensor] = None, return_attn: bool = False - ) -> Union[Tensor, Tuple[Tensor, Tensor]]: - """Compute the Nystrom attention.""" - _, n, _, h, m = *x.shape, self.num_heads, self.num_landmarks - if n % m != 0: - x, mask = self._pad_sequence(x, mask, n, m) - - q, k, v = self.qkv_fn(x).chunk(3, dim=-1) - q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v)) - - q = q * self.scale - - out, attn = self._nystrom_attention(q, k, v, mask, n, m, return_attn) - - # Add depth-wise convolutional residual of values - if self.residual is not None: - out += self.residual(out) - - out = rearrange(out, "b h n d -> b n (h d)", h=h) - out = self.fc_out(out) - out = out[:, -n:] - - if return_attn: - return out, attn - return out diff --git a/text_recognizer/networks/transformer/nystromer/nystromer.py b/text_recognizer/networks/transformer/nystromer/nystromer.py deleted file mode 100644 index 2113f1f..0000000 --- a/text_recognizer/networks/transformer/nystromer/nystromer.py +++ /dev/null @@ -1,64 +0,0 @@ -"""Nyströmer encoder. - -Stolen from: - https://github.com/lucidrains/nystrom-attention/blob/main/nystrom_attention/nystrom_attention.py - -""" -from typing import Optional - -from torch import nn, Tensor - -from text_recognizer.networks.transformer.mlp import FeedForward -from text_recognizer.networks.transformer.norm import PreNorm -from text_recognizer.networks.transformer.nystromer.attention import NystromAttention - - -class Nystromer(nn.Module): - def __init__( - self, - *, - dim: int, - depth: int, - dim_head: int = 64, - num_heads: int = 8, - num_landmarks: int = 256, - inverse_iter: int = 6, - residual: bool = True, - residual_conv_kernel: int = 33, - dropout_rate: float = 0.0, - glu: bool = True, - ) -> None: - super().__init__() - self.dim = dim - self.layers = nn.ModuleList( - [ - nn.ModuleList( - [ - PreNorm( - dim, - NystromAttention( - dim=dim, - dim_head=dim_head, - num_heads=num_heads, - num_landmarks=num_landmarks, - inverse_iter=inverse_iter, - residual=residual, - residual_conv_kernel=residual_conv_kernel, - dropout_rate=dropout_rate, - ), - ), - PreNorm( - dim, - FeedForward(dim=dim, glu=glu, dropout_rate=dropout_rate), - ), - ] - ) - for _ in range(depth) - ] - ) - - def forward(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor: - for attn, ff in self.layers: - x = attn(x, mask=mask) + x - x = ff(x) + x - return x |