diff options
Diffstat (limited to 'text_recognizer/network')
20 files changed, 662 insertions, 0 deletions
diff --git a/text_recognizer/network/__init__.py b/text_recognizer/network/__init__.py new file mode 100644 index 0000000..d9ef58b --- /dev/null +++ b/text_recognizer/network/__init__.py @@ -0,0 +1 @@ +"""Network modules""" diff --git a/text_recognizer/network/convnext/__init__.py b/text_recognizer/network/convnext/__init__.py new file mode 100644 index 0000000..dcff3fc --- /dev/null +++ b/text_recognizer/network/convnext/__init__.py @@ -0,0 +1,7 @@ +"""Convnext module.""" +from text_recognizer.network.convnext.attention import ( + Attention, + FeedForward, + TransformerBlock, +) +from text_recognizer.network.convnext.convnext import ConvNext diff --git a/text_recognizer/network/convnext/attention.py b/text_recognizer/network/convnext/attention.py new file mode 100644 index 0000000..6bc9692 --- /dev/null +++ b/text_recognizer/network/convnext/attention.py @@ -0,0 +1,79 @@ +"""Convolution self attention block.""" + +import torch.nn.functional as F +from einops import rearrange +from torch import Tensor, einsum, nn + +from text_recognizer.network.convnext.norm import LayerNorm +from text_recognizer.network.convnext.residual import Residual + + +def l2norm(t: Tensor) -> Tensor: + return F.normalize(t, dim=-1) + + +class FeedForward(nn.Module): + def __init__(self, dim: int, mult: int = 4) -> None: + super().__init__() + inner_dim = int(dim * mult) + self.fn = Residual( + nn.Sequential( + LayerNorm(dim), + nn.Conv2d(dim, inner_dim, 1, bias=False), + nn.GELU(), + LayerNorm(inner_dim), + nn.Conv2d(inner_dim, dim, 1, bias=False), + ) + ) + + def forward(self, x: Tensor) -> Tensor: + return self.fn(x) + + +class Attention(nn.Module): + def __init__( + self, dim: int, heads: int = 4, dim_head: int = 64, scale: int = 8 + ) -> None: + super().__init__() + self.scale = scale + self.heads = heads + inner_dim = heads * dim_head + self.norm = LayerNorm(dim) + + self.to_qkv = nn.Conv2d(dim, inner_dim * 3, 1, bias=False) + self.to_out = nn.Conv2d(inner_dim, dim, 1, bias=False) + + def forward(self, x: Tensor) -> Tensor: + h, w = x.shape[-2:] + + residual = x.clone() + + x = self.norm(x) + + q, k, v = self.to_qkv(x).chunk(3, dim=1) + q, k, v = map( + lambda t: rearrange(t, "b (h c) ... -> b h (...) c", h=self.heads), + (q, k, v), + ) + + q, k = map(l2norm, (q, k)) + + sim = einsum("b h i d, b h j d -> b h i j", q, k) * self.scale + attn = sim.softmax(dim=-1) + + out = einsum("b h i j, b h j d -> b h i d", attn, v) + + out = rearrange(out, "b h (x y) d -> b (h d) x y", x=h, y=w) + return self.to_out(out) + residual + + +class TransformerBlock(nn.Module): + def __init__(self, attn: Attention, ff: FeedForward) -> None: + super().__init__() + self.attn = attn + self.ff = ff + + def forward(self, x: Tensor) -> Tensor: + x = self.attn(x) + x = self.ff(x) + return x diff --git a/text_recognizer/network/convnext/convnext.py b/text_recognizer/network/convnext/convnext.py new file mode 100644 index 0000000..6acf059 --- /dev/null +++ b/text_recognizer/network/convnext/convnext.py @@ -0,0 +1,77 @@ +"""ConvNext module.""" +from typing import Optional, Sequence + +from torch import Tensor, nn + +from text_recognizer.network.convnext.attention import TransformerBlock +from text_recognizer.network.convnext.downsample import Downsample +from text_recognizer.network.convnext.norm import LayerNorm + + +class ConvNextBlock(nn.Module): + """ConvNext block.""" + + def __init__(self, dim: int, dim_out: int, mult: int) -> None: + super().__init__() + self.ds_conv = nn.Conv2d( + dim, dim, kernel_size=(7, 7), padding="same", groups=dim + ) + self.net = nn.Sequential( + LayerNorm(dim), + nn.Conv2d(dim, dim_out * mult, kernel_size=(3, 3), padding="same"), + nn.GELU(), + nn.Conv2d(dim_out * mult, dim_out, kernel_size=(3, 3), padding="same"), + ) + self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity() + + def forward(self, x: Tensor) -> Tensor: + h = self.ds_conv(x) + h = self.net(h) + return h + self.res_conv(x) + + +class ConvNext(nn.Module): + def __init__( + self, + dim: int = 16, + dim_mults: Sequence[int] = (2, 4, 8), + depths: Sequence[int] = (3, 3, 6), + downsampling_factors: Sequence[Sequence[int]] = ((2, 2), (2, 2), (2, 2)), + attn: Optional[TransformerBlock] = None, + ) -> None: + super().__init__() + dims = (dim, *map(lambda m: m * dim, dim_mults)) + self.attn = attn if attn is not None else nn.Identity() + self.out_channels = dims[-1] + self.stem = nn.Conv2d(1, dims[0], kernel_size=7, padding="same") + self.layers = nn.ModuleList([]) + + for i in range(len(dims) - 1): + dim_in, dim_out = dims[i], dims[i + 1] + self.layers.append( + nn.ModuleList( + [ + ConvNextBlock(dim_in, dim_in, 2), + nn.ModuleList( + [ConvNextBlock(dim_in, dim_in, 2) for _ in range(depths[i])] + ), + Downsample(dim_in, dim_out, downsampling_factors[i]), + ] + ) + ) + self.norm = LayerNorm(dims[-1]) + + def _init_weights(self, m): + if isinstance(m, (nn.Conv2d, nn.Linear)): + nn.init.trunc_normal_(m.weight, std=0.02) + nn.init.constant_(m.bias, 0) + + def forward(self, x: Tensor) -> Tensor: + x = self.stem(x) + for init_block, blocks, down in self.layers: + x = init_block(x) + for fn in blocks: + x = fn(x) + x = down(x) + x = self.attn(x) + return self.norm(x) diff --git a/text_recognizer/network/convnext/downsample.py b/text_recognizer/network/convnext/downsample.py new file mode 100644 index 0000000..a8a0466 --- /dev/null +++ b/text_recognizer/network/convnext/downsample.py @@ -0,0 +1,21 @@ +"""Convnext downsample module.""" +from typing import Tuple + +from einops.layers.torch import Rearrange +from torch import Tensor, nn + + +class Downsample(nn.Module): + """Downsamples feature maps by patches.""" + + def __init__(self, dim: int, dim_out: int, factors: Tuple[int, int]) -> None: + super().__init__() + s1, s2 = factors + self.fn = nn.Sequential( + Rearrange("b c (h s1) (w s2) -> b (c s1 s2) h w", s1=s1, s2=s2), + nn.Conv2d(dim * s1 * s2, dim_out, 1), + ) + + def forward(self, x: Tensor) -> Tensor: + """Applies patch function.""" + return self.fn(x) diff --git a/text_recognizer/network/convnext/norm.py b/text_recognizer/network/convnext/norm.py new file mode 100644 index 0000000..3355de9 --- /dev/null +++ b/text_recognizer/network/convnext/norm.py @@ -0,0 +1,18 @@ +"""Layer norm for conv layers.""" +import torch +from torch import Tensor, nn + + +class LayerNorm(nn.Module): + """Layer norm for convolutions.""" + + def __init__(self, dim: int) -> None: + super().__init__() + self.gamma = nn.Parameter(torch.ones(1, dim, 1, 1)) + + def forward(self, x: Tensor) -> Tensor: + """Applies layer norm.""" + eps = 1e-5 if x.dtype == torch.float32 else 1e-3 + var = torch.var(x, dim=1, unbiased=False, keepdim=True) + mean = torch.mean(x, dim=1, keepdim=True) + return (x - mean) / (var + eps).sqrt() * self.gamma diff --git a/text_recognizer/network/convnext/residual.py b/text_recognizer/network/convnext/residual.py new file mode 100644 index 0000000..dfc2847 --- /dev/null +++ b/text_recognizer/network/convnext/residual.py @@ -0,0 +1,16 @@ +"""Generic residual layer.""" +from typing import Callable + +from torch import Tensor, nn + + +class Residual(nn.Module): + """Residual layer.""" + + def __init__(self, fn: Callable) -> None: + super().__init__() + self.fn = fn + + def forward(self, x: Tensor) -> Tensor: + """Applies residual fn.""" + return self.fn(x) + x diff --git a/text_recognizer/network/transformer/__init__.py b/text_recognizer/network/transformer/__init__.py new file mode 100644 index 0000000..a3f3011 --- /dev/null +++ b/text_recognizer/network/transformer/__init__.py @@ -0,0 +1 @@ +"""Transformer modules.""" diff --git a/text_recognizer/network/transformer/attend.py b/text_recognizer/network/transformer/attend.py new file mode 100644 index 0000000..4e643fb --- /dev/null +++ b/text_recognizer/network/transformer/attend.py @@ -0,0 +1,94 @@ +from typing import Optional +from collections import namedtuple + +import torch +from torch import Tensor, einsum, nn +from einops import rearrange +import torch.nn.functional as F + +Config = namedtuple( + "FlashAttentionConfig", ["enable_flash", "enable_math", "enable_mem_efficient"] +) + + +class Attend(nn.Module): + def __init__(self, use_flash: bool) -> None: + super().__init__() + self.use_flash = use_flash + self.cpu_cfg = Config(True, True, True) + self.cuda_cfg = None + if not torch.cuda.is_available(): + return + + device_properties = torch.cuda.get_device_properties(torch.device("cuda")) + if device_properties.major == 8 and device_properties.minor == 0: + self.cuda_cfg = Config(True, False, False) + else: + self.cuda_cfg = Config(False, True, True) + + def flash_attn(self, q: Tensor, k: Tensor, v: Tensor, causal: bool) -> Tensor: + cfg = self.cuda_cfg if q.is_cuda else self.cpu_cfg + with torch.backends.cuda.sdp_kernel(**cfg._asdict()): + out = F.scaled_dot_product_attention(q, k, v, is_causal=causal) + return out + + def atten( + self, + q: Tensor, + k: Tensor, + v: Tensor, + causal: bool, + mask: Optional[Tensor] = None, + ) -> Tensor: + b = q.shape[0] + energy = einsum("b h i d, b h j d -> b h i j", q, k) * self.scale + + mask_value = -torch.finfo(energy.dtype).max + + if mask is not None: + energy = apply_input_mask(b, k, energy, mask, mask_value) + + if causal: + energy = apply_causal_mask(energy, mask_value) + + attn = F.softmax(energy, dim=-1) + attn = self.dropout(attn) + return einsum("b h i j, b h j d -> b h i d", attn, v) + + def forward( + self, + q: Tensor, + k: Tensor, + v: Tensor, + causal: bool, + mask: Optional[Tensor] = None, + ) -> Tensor: + if self.use_flash: + return self.flash_attn(q, k, v, causal) + else: + return self.atten(q, k, v, causal, mask) + + +def apply_input_mask( + b: int, + k: Tensor, + energy: Tensor, + mask: Optional[Tensor], + mask_value: float, +) -> Tensor: + """Applies an input mask.""" + k_mask = torch.ones((b, k.shape[-2]), device=energy.device).bool() + q_mask = rearrange(mask, "b i -> b () i ()") + k_mask = rearrange(k_mask, "b j -> b () () j") + input_mask = q_mask * k_mask + return energy.masked_fill_(~input_mask, mask_value) + + +def apply_causal_mask( + energy: Tensor, + mask_value: float, +) -> Tensor: + """Applies a causal mask to the energy tensor.""" + i, j, device = *energy.shape[-2:], energy.device + causal_mask = torch.ones((i, j), device=device, dtype=torch.bool).triu(j - i + 1) + return energy.masked_fill(causal_mask, mask_value) diff --git a/text_recognizer/network/transformer/attention.py b/text_recognizer/network/transformer/attention.py new file mode 100644 index 0000000..8e18f8a --- /dev/null +++ b/text_recognizer/network/transformer/attention.py @@ -0,0 +1,56 @@ +"""Implements the attention module for the transformer.""" +from typing import Optional +from text_recognizer.network.transformer.norm import RMSNorm +from text_recognizer.network.transformer.attend import Attend + +import torch +from einops import rearrange +from torch import Tensor, nn + + +class Attention(nn.Module): + """Standard attention.""" + + def __init__( + self, + dim: int, + heads: int, + causal: bool = False, + dim_head: int = 64, + dropout_rate: float = 0.0, + use_flash: bool = True, + ) -> None: + super().__init__() + self.heads = heads + inner_dim = dim_head * heads + self.norm = nn.LayerNorm(dim) + self.to_q = nn.Linear(dim, inner_dim, bias=False) + self.to_k = nn.Linear(dim, inner_dim, bias=False) + self.to_v = nn.Linear(dim, inner_dim, bias=False) + # self.q_norm = RMSNorm(heads, dim_head) + # self.k_norm = RMSNorm(heads, dim_head) + self.attend = Attend(use_flash) + self.to_out = nn.Linear(inner_dim, dim, bias=False) + self.scale = dim**-0.5 + self.causal = causal + self.dropout_rate = dropout_rate + self.dropout = nn.Dropout(p=self.dropout_rate) + + def forward( + self, + x: Tensor, + context: Optional[Tensor] = None, + mask: Optional[Tensor] = None, + ) -> Tensor: + """Computes the attention.""" + x = self.norm(x) + q = self.to_q(x) + k = self.to_k(x if context is None else context) + v = self.to_v(x if context is None else context) + q, k, v = map( + lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.heads), (q, k, v) + ) + out = self.attend(q, k, v, self.causal, mask) + out = rearrange(out, "b h n d -> b n (h d)") + out = self.to_out(out) + return out diff --git a/text_recognizer/network/transformer/decoder.py b/text_recognizer/network/transformer/decoder.py new file mode 100644 index 0000000..06925ba --- /dev/null +++ b/text_recognizer/network/transformer/decoder.py @@ -0,0 +1,57 @@ +"""Transformer decoder module.""" +from typing import Optional +from torch import Tensor, nn + +from text_recognizer.network.transformer.attention import Attention +from text_recognizer.network.transformer.ff import FeedForward + + +class Decoder(nn.Module): + def __init__( + self, + dim: int, + inner_dim: int, + heads: int, + dim_head: int, + depth: int, + dropout_rate: float = 0.0, + ) -> None: + super().__init__() + self.norm = nn.LayerNorm(dim) + self.layers = nn.ModuleList( + [ + nn.ModuleList( + [ + Attention( + dim, + heads, + True, + dim_head, + dropout_rate, + ), + FeedForward(dim, inner_dim, dropout_rate), + Attention( + dim, + heads, + False, + dim_head, + dropout_rate, + ), + ] + ) + for _ in range(depth) + ] + ) + + def forward( + self, + x: Tensor, + context: Tensor, + mask: Optional[Tensor] = None, + ) -> Tensor: + """Applies decoder block on input signals.""" + for self_attn, ff, cross_attn in self.layers: + x = x + self_attn(x, mask=mask) + x = x + ff(x) + x = x + cross_attn(x, context=context) + return self.norm(x) diff --git a/text_recognizer/network/transformer/embedding/__init__.py b/text_recognizer/network/transformer/embedding/__init__.py new file mode 100644 index 0000000..bb3f904 --- /dev/null +++ b/text_recognizer/network/transformer/embedding/__init__.py @@ -0,0 +1 @@ +"""Positional encodings for transformers.""" diff --git a/text_recognizer/network/transformer/embedding/absolute.py b/text_recognizer/network/transformer/embedding/absolute.py new file mode 100644 index 0000000..08b2c2a --- /dev/null +++ b/text_recognizer/network/transformer/embedding/absolute.py @@ -0,0 +1,28 @@ +from typing import Optional + +import torch +from torch import nn, Tensor +from text_recognizer.network.transformer.embedding.l2_norm import l2_norm + + +class AbsolutePositionalEmbedding(nn.Module): + def __init__(self, dim: int, max_length: int, use_l2: bool = False) -> None: + super().__init__() + self.scale = dim**-0.5 if not use_l2 else 1.0 + self.max_length = max_length + self.use_l2 = use_l2 + self.to_embedding = nn.Embedding(max_length, dim) + if self.use_l2: + nn.init.normal_(self.to_embedding.weight, std=1e-5) + + def forward(self, x: Tensor, pos: Optional[Tensor] = None) -> Tensor: + n, device = x.shape[1], x.device + assert ( + n <= self.max_length + ), f"Sequence length {n} is greater than the maximum positional embedding {self.max_length}" + + if pos is None: + pos = torch.arange(n, device=device) + + embedding = self.to_embedding(pos) * self.scale + return l2_norm(embedding) if self.use_l2 else embedding diff --git a/text_recognizer/network/transformer/embedding/l2_norm.py b/text_recognizer/network/transformer/embedding/l2_norm.py new file mode 100644 index 0000000..0e48bca --- /dev/null +++ b/text_recognizer/network/transformer/embedding/l2_norm.py @@ -0,0 +1,9 @@ +from einops import rearrange +import torch.nn.functional as F +from torch import Tensor + + +def l2_norm(t: Tensor, groups=1) -> Tensor: + t = rearrange(t, "... (g d) -> ... g d", g=groups) + t = F.normalize(t, p=2, dim=-1) + return rearrange(t, "... g d -> ... (g d)") diff --git a/text_recognizer/network/transformer/embedding/sincos.py b/text_recognizer/network/transformer/embedding/sincos.py new file mode 100644 index 0000000..ed6b0ab --- /dev/null +++ b/text_recognizer/network/transformer/embedding/sincos.py @@ -0,0 +1,13 @@ +import torch + + +def sincos_2d(h, w, dim, temperature: int = 10000, dtype=torch.float32): + y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij") + assert (dim % 4) == 0, "feature dimension must be multiple of 4 for sincos emb" + omega = torch.arange(dim // 4) / (dim // 4 - 1) + omega = 1.0 / (temperature**omega) + + y = y.flatten()[:, None] * omega[None, :] + x = x.flatten()[:, None] * omega[None, :] + pe = torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim=1) + return pe.type(dtype) diff --git a/text_recognizer/network/transformer/embedding/token.py b/text_recognizer/network/transformer/embedding/token.py new file mode 100644 index 0000000..1df2fd6 --- /dev/null +++ b/text_recognizer/network/transformer/embedding/token.py @@ -0,0 +1,18 @@ +from torch import nn, Tensor + +from text_recognizer.network.transformer.embedding.l2_norm import l2_norm + + +class TokenEmbedding(nn.Module): + def __init__(self, num_tokens: int, dim: int, use_l2: bool = True) -> None: + super().__init__() + self.use_l2 = use_l2 + self.to_embedding = nn.Embedding(num_tokens, dim) + if self.use_l2: + nn.init.normal_(self.to_embedding.weight, std=1e-5) + else: + nn.init.kaiming_normal_(self.to_embedding.weight) + + def forward(self, x: Tensor) -> Tensor: + embedding = self.to_embedding(x) + return l2_norm(embedding) if self.use_l2 else embedding diff --git a/text_recognizer/network/transformer/encoder.py b/text_recognizer/network/transformer/encoder.py new file mode 100644 index 0000000..ea4b0b3 --- /dev/null +++ b/text_recognizer/network/transformer/encoder.py @@ -0,0 +1,46 @@ +"""Transformer encoder module.""" +from torch import Tensor, nn + +from text_recognizer.network.transformer.attention import Attention +from text_recognizer.network.transformer.ff import FeedForward + + +class Encoder(nn.Module): + def __init__( + self, + dim: int, + inner_dim: int, + heads: int, + dim_head: int, + depth: int, + dropout_rate: float = 0.0, + ) -> None: + super().__init__() + self.norm = nn.LayerNorm(dim) + self.layers = nn.ModuleList( + [ + nn.ModuleList( + [ + Attention( + dim, + heads, + False, + dim_head, + dropout_rate, + ), + FeedForward(dim, inner_dim, dropout_rate), + ] + ) + for _ in range(depth) + ] + ) + + def forward( + self, + x: Tensor, + ) -> Tensor: + """Applies decoder block on input signals.""" + for self_attn, ff in self.layers: + x = x + self_attn(x) + x = x + ff(x) + return self.norm(x) diff --git a/text_recognizer/network/transformer/ff.py b/text_recognizer/network/transformer/ff.py new file mode 100644 index 0000000..9181323 --- /dev/null +++ b/text_recognizer/network/transformer/ff.py @@ -0,0 +1,22 @@ +"""Feedforward layer in transformer.""" +from torch import Tensor, nn + + +class FeedForward(nn.Module): + def __init__( + self, + dim: int, + inner_dim: int, + dropout_rate: float = 0.0, + ) -> None: + super().__init__() + self.ff = nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, inner_dim), + nn.GELU(), + nn.Dropout(dropout_rate), + nn.Linear(inner_dim, dim), + ) + + def forward(self, x: Tensor) -> Tensor: + return self.ff(x) diff --git a/text_recognizer/network/transformer/norm.py b/text_recognizer/network/transformer/norm.py new file mode 100644 index 0000000..2737754 --- /dev/null +++ b/text_recognizer/network/transformer/norm.py @@ -0,0 +1,22 @@ +"""Normalization layers for transformers. + +Copied from lucidrains: + https://github.com/lucidrains/x-transformers/blob/main/x_transformers/x_transformers.py + +""" +import torch +from torch import Tensor, nn +import torch.nn.functional as F + + +class RMSNorm(nn.Module): + """Root mean square layer normalization.""" + + def __init__(self, heads: int, dim: int) -> None: + super().__init__() + self.scale = dim**-0.5 + self.gamma = nn.Parameter(torch.ones(heads, 1, dim)) + + def forward(self, x: Tensor) -> Tensor: + """Applies normalization.""" + return F.normalize(x, dim=-1) * self.scale * self.gamma diff --git a/text_recognizer/network/vit.py b/text_recognizer/network/vit.py new file mode 100644 index 0000000..80176a8 --- /dev/null +++ b/text_recognizer/network/vit.py @@ -0,0 +1,76 @@ +"""Transformer module.""" +from typing import Type + +from einops.layers.torch import Rearrange +from torch import Tensor, nn + +from text_recognizer.network.transformer.embedding.token import TokenEmbedding +from text_recognizer.network.transformer.embedding.sincos import sincos_2d +from text_recognizer.network.transformer.decoder import Decoder +from text_recognizer.network.transformer.encoder import Encoder + + +class VisionTransformer(nn.Module): + def __init__( + self, + image_height: int, + image_width: int, + patch_height: int, + patch_width: int, + dim: int, + num_classes: int, + encoder: Encoder, + decoder: Decoder, + token_embedding: TokenEmbedding, + pos_embedding: Type[nn.Module], + tie_embeddings: bool, + pad_index: int, + ) -> None: + super().__init__() + patch_dim = patch_height * patch_width + self.to_patch_embedding = nn.Sequential( + Rearrange( + "b c (h ph) (w pw) -> b (h w) (ph pw c)", + ph=patch_height, + pw=patch_width, + ), + nn.LayerNorm(patch_dim), + nn.Linear(patch_dim, dim), + nn.LayerNorm(dim), + ) + self.patch_embedding = sincos_2d( + h=image_height // patch_height, w=image_width // patch_width, dim=dim + ) + self.pos_embedding = pos_embedding + self.token_embedding = token_embedding + self.to_logits = ( + nn.Linear(dim, num_classes) + if not tie_embeddings + else lambda t: t @ self.token_embedding.to_embedding.weight.t() + ) + self.encoder = encoder + self.decoder = decoder + self.pad_index = pad_index + + def encode(self, img: Tensor) -> Tensor: + x = self.to_patch_embedding(img) + x += self.patch_embedding.to(img.device, dtype=img.dtype) + return self.encoder(x) + + def decode(self, text: Tensor, context: Tensor) -> Tensor: + text = text.long() + mask = text != self.pad_index + tokens = self.token_embedding(text) + tokens = tokens + self.pos_embedding(tokens) + output = self.decoder(tokens, context) + return self.to_logits(output) + + def forward( + self, + img: Tensor, + text: Tensor, + ) -> Tensor: + """Applies decoder block on input signals.""" + context = self.encode(img) + logits = self.decode(text, context) + return logits.permute(0, 2, 1) |