From 49ca6ade1a19f7f9c702171537fe4be0dfcda66d Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Fri, 25 Aug 2023 23:19:14 +0200
Subject: Rename and add flash atten

---
 text_recognizer/data/base_data_module.py           |   4 +-
 text_recognizer/network/__init__.py                |   1 +
 text_recognizer/network/convnext/__init__.py       |   7 ++
 text_recognizer/network/convnext/attention.py      |  79 +++++++++++++
 text_recognizer/network/convnext/convnext.py       |  77 ++++++++++++
 text_recognizer/network/convnext/downsample.py     |  21 ++++
 text_recognizer/network/convnext/norm.py           |  18 +++
 text_recognizer/network/convnext/residual.py       |  16 +++
 text_recognizer/network/transformer/__init__.py    |   1 +
 text_recognizer/network/transformer/attend.py      |  94 +++++++++++++++
 text_recognizer/network/transformer/attention.py   |  56 +++++++++
 text_recognizer/network/transformer/decoder.py     |  57 +++++++++
 .../network/transformer/embedding/__init__.py      |   1 +
 .../network/transformer/embedding/absolute.py      |  28 +++++
 .../network/transformer/embedding/l2_norm.py       |   9 ++
 .../network/transformer/embedding/sincos.py        |  13 +++
 .../network/transformer/embedding/token.py         |  18 +++
 text_recognizer/network/transformer/encoder.py     |  46 ++++++++
 text_recognizer/network/transformer/ff.py          |  22 ++++
 text_recognizer/network/transformer/norm.py        |  22 ++++
 text_recognizer/network/vit.py                     |  76 ++++++++++++
 text_recognizer/networks/__init__.py               |   2 -
 text_recognizer/networks/conv_transformer.py       |  49 --------
 text_recognizer/networks/convnext/__init__.py      |   7 --
 text_recognizer/networks/convnext/attention.py     |  79 -------------
 text_recognizer/networks/convnext/convnext.py      |  77 ------------
 text_recognizer/networks/convnext/downsample.py    |  21 ----
 text_recognizer/networks/convnext/norm.py          |  18 ---
 text_recognizer/networks/convnext/residual.py      |  16 ---
 text_recognizer/networks/image_encoder.py          |  45 -------
 text_recognizer/networks/text_decoder.py           |  55 ---------
 text_recognizer/networks/transformer/__init__.py   |   6 -
 text_recognizer/networks/transformer/attention.py  | 109 -----------------
 text_recognizer/networks/transformer/decoder.py    |  41 -------
 .../networks/transformer/decoder_block.py          |  44 -------
 .../networks/transformer/embeddings/__init__.py    |   1 -
 .../networks/transformer/embeddings/axial.py       | 104 -----------------
 .../networks/transformer/embeddings/rotary.py      |  67 -----------
 text_recognizer/networks/transformer/ff.py         |  45 -------
 text_recognizer/networks/transformer/norm.py       |  51 --------
 text_recognizer/optimizer/__init__.py              |   0
 text_recognizer/optimizer/laprop.py                | 129 +++++++++++++++++++++
 text_recognizer/optimizers/__init__.py             |   0
 text_recognizer/optimizers/laprop.py               | 129 ---------------------
 44 files changed, 793 insertions(+), 968 deletions(-)
 create mode 100644 text_recognizer/network/__init__.py
 create mode 100644 text_recognizer/network/convnext/__init__.py
 create mode 100644 text_recognizer/network/convnext/attention.py
 create mode 100644 text_recognizer/network/convnext/convnext.py
 create mode 100644 text_recognizer/network/convnext/downsample.py
 create mode 100644 text_recognizer/network/convnext/norm.py
 create mode 100644 text_recognizer/network/convnext/residual.py
 create mode 100644 text_recognizer/network/transformer/__init__.py
 create mode 100644 text_recognizer/network/transformer/attend.py
 create mode 100644 text_recognizer/network/transformer/attention.py
 create mode 100644 text_recognizer/network/transformer/decoder.py
 create mode 100644 text_recognizer/network/transformer/embedding/__init__.py
 create mode 100644 text_recognizer/network/transformer/embedding/absolute.py
 create mode 100644 text_recognizer/network/transformer/embedding/l2_norm.py
 create mode 100644 text_recognizer/network/transformer/embedding/sincos.py
 create mode 100644 text_recognizer/network/transformer/embedding/token.py
 create mode 100644 text_recognizer/network/transformer/encoder.py
 create mode 100644 text_recognizer/network/transformer/ff.py
 create mode 100644 text_recognizer/network/transformer/norm.py
 create mode 100644 text_recognizer/network/vit.py
 delete mode 100644 text_recognizer/networks/__init__.py
 delete mode 100644 text_recognizer/networks/conv_transformer.py
 delete mode 100644 text_recognizer/networks/convnext/__init__.py
 delete mode 100644 text_recognizer/networks/convnext/attention.py
 delete mode 100644 text_recognizer/networks/convnext/convnext.py
 delete mode 100644 text_recognizer/networks/convnext/downsample.py
 delete mode 100644 text_recognizer/networks/convnext/norm.py
 delete mode 100644 text_recognizer/networks/convnext/residual.py
 delete mode 100644 text_recognizer/networks/image_encoder.py
 delete mode 100644 text_recognizer/networks/text_decoder.py
 delete mode 100644 text_recognizer/networks/transformer/__init__.py
 delete mode 100644 text_recognizer/networks/transformer/attention.py
 delete mode 100644 text_recognizer/networks/transformer/decoder.py
 delete mode 100644 text_recognizer/networks/transformer/decoder_block.py
 delete mode 100644 text_recognizer/networks/transformer/embeddings/__init__.py
 delete mode 100644 text_recognizer/networks/transformer/embeddings/axial.py
 delete mode 100644 text_recognizer/networks/transformer/embeddings/rotary.py
 delete mode 100644 text_recognizer/networks/transformer/ff.py
 delete mode 100644 text_recognizer/networks/transformer/norm.py
 create mode 100644 text_recognizer/optimizer/__init__.py
 create mode 100644 text_recognizer/optimizer/laprop.py
 delete mode 100644 text_recognizer/optimizers/__init__.py
 delete mode 100644 text_recognizer/optimizers/laprop.py

diff --git a/text_recognizer/data/base_data_module.py b/text_recognizer/data/base_data_module.py
index bd6fd99..3cb0aa2 100644
--- a/text_recognizer/data/base_data_module.py
+++ b/text_recognizer/data/base_data_module.py
@@ -2,7 +2,7 @@
 from pathlib import Path
 from typing import Callable, Dict, Optional, Tuple, TypeVar
 
-from pytorch_lightning import LightningDataModule
+import pytorch_lightning as L
 from torch.utils.data import DataLoader
 
 from text_recognizer.data.base_dataset import BaseDataset
@@ -19,7 +19,7 @@ def load_and_print_info(data_module_class: type) -> None:
     print(dataset)
 
 
-class BaseDataModule(LightningDataModule):
+class BaseDataModule(L.LightningDataModule):
     """Base PyTorch Lightning DataModule."""
 
     def __init__(
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)
diff --git a/text_recognizer/networks/__init__.py b/text_recognizer/networks/__init__.py
deleted file mode 100644
index f921882..0000000
--- a/text_recognizer/networks/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-"""Network modules"""
-from text_recognizer.networks.conv_transformer import ConvTransformer
diff --git a/text_recognizer/networks/conv_transformer.py b/text_recognizer/networks/conv_transformer.py
deleted file mode 100644
index d36162a..0000000
--- a/text_recognizer/networks/conv_transformer.py
+++ /dev/null
@@ -1,49 +0,0 @@
-"""Base network module."""
-from typing import Type
-
-from torch import Tensor, nn
-
-from text_recognizer.networks.transformer.decoder import Decoder
-
-
-class ConvTransformer(nn.Module):
-    """Base transformer network."""
-
-    def __init__(
-        self,
-        encoder: Type[nn.Module],
-        decoder: Decoder,
-    ) -> None:
-        super().__init__()
-        self.encoder = encoder
-        self.decoder = decoder
-
-    def encode(self, img: Tensor) -> Tensor:
-        """Encodes images to latent representation."""
-        return self.encoder(img)
-
-    def decode(self, tokens: Tensor, img_features: Tensor) -> Tensor:
-        """Decodes latent images embedding into characters."""
-        return self.decoder(tokens, img_features)
-
-    def forward(self, img: Tensor, tokens: Tensor) -> Tensor:
-        """Encodes images into token logtis.
-
-        Args:
-            img (Tensor): Input image(s).
-            tokens (Tensor): token embeddings.
-
-        Shapes:
-            - img: :math: `(B, 1, H, W)`
-            - tokens: :math: `(B, Sy)`
-            - logits: :math: `(B, Sy, C)`
-
-            where B is the batch size, H is the image height, W is the image
-            width, Sy the output length, and C is the number of classes.
-
-        Returns:
-            Tensor: Sequence of logits.
-        """
-        img_features = self.encode(img)
-        logits = self.decode(tokens, img_features)
-        return logits
diff --git a/text_recognizer/networks/convnext/__init__.py b/text_recognizer/networks/convnext/__init__.py
deleted file mode 100644
index faebe6f..0000000
--- a/text_recognizer/networks/convnext/__init__.py
+++ /dev/null
@@ -1,7 +0,0 @@
-"""Convnext module."""
-from text_recognizer.networks.convnext.attention import (
-    Attention,
-    FeedForward,
-    TransformerBlock,
-)
-from text_recognizer.networks.convnext.convnext import ConvNext
diff --git a/text_recognizer/networks/convnext/attention.py b/text_recognizer/networks/convnext/attention.py
deleted file mode 100644
index 1334feb..0000000
--- a/text_recognizer/networks/convnext/attention.py
+++ /dev/null
@@ -1,79 +0,0 @@
-"""Convolution self attention block."""
-
-import torch.nn.functional as F
-from einops import rearrange
-from torch import Tensor, einsum, nn
-
-from text_recognizer.networks.convnext.norm import LayerNorm
-from text_recognizer.networks.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/networks/convnext/convnext.py b/text_recognizer/networks/convnext/convnext.py
deleted file mode 100644
index 9419a15..0000000
--- a/text_recognizer/networks/convnext/convnext.py
+++ /dev/null
@@ -1,77 +0,0 @@
-"""ConvNext module."""
-from typing import Optional, Sequence
-
-from torch import Tensor, nn
-
-from text_recognizer.networks.convnext.attention import TransformerBlock
-from text_recognizer.networks.convnext.downsample import Downsample
-from text_recognizer.networks.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/networks/convnext/downsample.py b/text_recognizer/networks/convnext/downsample.py
deleted file mode 100644
index a8a0466..0000000
--- a/text_recognizer/networks/convnext/downsample.py
+++ /dev/null
@@ -1,21 +0,0 @@
-"""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/networks/convnext/norm.py b/text_recognizer/networks/convnext/norm.py
deleted file mode 100644
index 3355de9..0000000
--- a/text_recognizer/networks/convnext/norm.py
+++ /dev/null
@@ -1,18 +0,0 @@
-"""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/networks/convnext/residual.py b/text_recognizer/networks/convnext/residual.py
deleted file mode 100644
index dfc2847..0000000
--- a/text_recognizer/networks/convnext/residual.py
+++ /dev/null
@@ -1,16 +0,0 @@
-"""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/networks/image_encoder.py b/text_recognizer/networks/image_encoder.py
deleted file mode 100644
index ab60560..0000000
--- a/text_recognizer/networks/image_encoder.py
+++ /dev/null
@@ -1,45 +0,0 @@
-"""Encodes images to latent embeddings."""
-from typing import Tuple, Type
-
-from torch import Tensor, nn
-
-from text_recognizer.networks.transformer.embeddings.axial import (
-    AxialPositionalEmbeddingImage,
-)
-
-
-class ImageEncoder(nn.Module):
-    """Encodes images to latent embeddings."""
-
-    def __init__(
-        self,
-        encoder: Type[nn.Module],
-        pixel_embedding: AxialPositionalEmbeddingImage,
-    ) -> None:
-        super().__init__()
-        self.encoder = encoder
-        self.pixel_embedding = pixel_embedding
-
-    def forward(self, img: Tensor) -> Tensor:
-        """Encodes an image into a latent feature vector.
-
-        Args:
-            img (Tensor): Image tensor.
-
-        Shape:
-            - x: :math: `(B, C, H, W)`
-            - z: :math: `(B, Sx, D)`
-
-            where Sx is the length of the flattened feature maps projected from
-            the encoder. D latent dimension for each pixel in the projected
-            feature maps.
-
-        Returns:
-            Tensor: A Latent embedding of the image.
-        """
-        z = self.encoder(img)
-        z = z + self.pixel_embedding(z)
-        z = z.flatten(start_dim=2)
-        # Permute tensor from [B, E, Ho * Wo] to [B, Sx, E]
-        z = z.permute(0, 2, 1)
-        return z
diff --git a/text_recognizer/networks/text_decoder.py b/text_recognizer/networks/text_decoder.py
deleted file mode 100644
index 500bcf9..0000000
--- a/text_recognizer/networks/text_decoder.py
+++ /dev/null
@@ -1,55 +0,0 @@
-"""Text decoder."""
-import torch
-from torch import Tensor, nn
-
-from text_recognizer.networks.transformer.decoder import Decoder
-
-
-class TextDecoder(nn.Module):
-    """Decodes images to token logits."""
-
-    def __init__(
-        self,
-        dim: int,
-        num_classes: int,
-        pad_index: Tensor,
-        decoder: Decoder,
-    ) -> None:
-        super().__init__()
-        self.dim = dim
-        self.num_classes = num_classes
-        self.pad_index = pad_index
-        self.decoder = decoder
-        self.token_embedding = nn.Embedding(
-            num_embeddings=self.num_classes, embedding_dim=self.dim
-        )
-        self.to_logits = nn.Linear(in_features=self.dim, out_features=self.num_classes)
-
-    def forward(self, tokens: Tensor, img_features: Tensor) -> Tensor:
-        """Decodes latent images embedding into logit tokens.
-
-        Args:
-            tokens (Tensor): Token indecies.
-            img_features (Tensor): Latent images embedding.
-
-        Shapes:
-            - tokens: :math: `(B, Sy)`
-            - img_features: :math: `(B, Sx, D)`
-            - logits: :math: `(B, Sy, C)`
-
-            where Sy is the length of the output, C is the number of classes
-            and D is the hidden dimension.
-
-        Returns:
-            Tensor: Sequence of logits.
-        """
-        tokens = tokens.long()
-        mask = tokens != self.pad_index
-        tokens = self.token_embedding(tokens)
-        tokens = self.decoder(x=tokens, context=img_features, mask=mask)
-        logits = (
-            tokens @ torch.transpose(self.token_embedding.weight.to(tokens.dtype), 0, 1)
-        ).float()
-        logits = self.to_logits(tokens)  # [B, Sy, C]
-        logits = logits.permute(0, 2, 1)  # [B, C, Sy]
-        return logits
diff --git a/text_recognizer/networks/transformer/__init__.py b/text_recognizer/networks/transformer/__init__.py
deleted file mode 100644
index 0d17deb..0000000
--- a/text_recognizer/networks/transformer/__init__.py
+++ /dev/null
@@ -1,6 +0,0 @@
-"""Transformer modules."""
-from text_recognizer.networks.transformer.attention import Attention
-from text_recognizer.networks.transformer.decoder import Decoder, DecoderBlock
-from text_recognizer.networks.transformer.embeddings.rotary import RotaryEmbedding
-from text_recognizer.networks.transformer.ff import FeedForward
-from text_recognizer.networks.transformer.norm import RMSNorm
diff --git a/text_recognizer/networks/transformer/attention.py b/text_recognizer/networks/transformer/attention.py
deleted file mode 100644
index 85f513e..0000000
--- a/text_recognizer/networks/transformer/attention.py
+++ /dev/null
@@ -1,109 +0,0 @@
-"""Implementes the attention module for the transformer."""
-from typing import Optional
-
-import torch
-import torch.nn.functional as F
-from einops import rearrange
-from torch import Tensor, einsum, nn
-
-from text_recognizer.networks.transformer.embeddings.rotary import (
-    RotaryEmbedding,
-)
-
-
-class Attention(nn.Module):
-    """Standard attention."""
-
-    def __init__(
-        self,
-        dim: int,
-        num_heads: int,
-        causal: bool = False,
-        dim_head: int = 64,
-        dropout_rate: float = 0.0,
-    ) -> None:
-        super().__init__()
-        self.dim = dim
-        self.scale = self.dim**-0.5
-        self.num_heads = num_heads
-        self.dim_head = dim_head
-
-        self.causal = causal
-        self.dropout_rate = dropout_rate
-
-        # Single key/value head
-        k_dim = dim_head
-        v_dim = dim_head
-
-        out_dim = self.num_heads * self.dim_head
-
-        self.to_q = nn.Linear(self.dim, out_dim, bias=False)
-        self.to_k = nn.Linear(self.dim, k_dim, bias=False)
-        self.to_v = nn.Linear(self.dim, v_dim, bias=False)
-
-        self.dropout = nn.Dropout(p=self.dropout_rate)
-
-        # Feedforward
-        self.fc = nn.Linear(out_dim, self.dim)
-
-    def forward(
-        self,
-        x: Tensor,
-        context: Optional[Tensor] = None,
-        mask: Optional[Tensor] = None,
-        rotary_embedding: Optional[RotaryEmbedding] = None,
-    ) -> Tensor:
-        """Computes the attention."""
-        b, device = x.shape[0], x.device
-
-        q = self.to_q(x)
-        q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads)
-        k = self.to_k(context) if context is not None else self.to_k(x)
-        v = self.to_v(context) if context is not None else self.to_v(x)
-
-        if rotary_embedding is not None:
-            q, k, v = map(lambda t: rotary_embedding.rotate(t), (q, k, v))
-
-        energy = einsum("b h i d, b j d -> b h i j", q, k) * self.scale
-        mask_value = -torch.finfo(energy.dtype).max
-        energy = apply_input_mask(b, k, energy, mask, mask_value, device)
-        if self.causal:
-            energy = apply_causal_mask(energy, mask, mask_value, device)
-
-        attn = F.softmax(energy, dim=-1)
-        attn = self.dropout(attn)
-        out = einsum("b h i j, b j d -> b h i d", attn, v)
-        out = rearrange(out, "b h n d -> b n (h d)")
-        out = self.fc(out)
-        return out
-
-
-def apply_input_mask(
-    b: int,
-    k: Tensor,
-    energy: Tensor,
-    mask: Optional[Tensor],
-    mask_value: Tensor,
-    device: str,
-) -> Tensor:
-    """Applies an input mask."""
-    if mask is not None:
-        k_mask = torch.ones((b, k.shape[-2]), device=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
-
-        energy = energy.masked_fill_(~input_mask, mask_value)
-    return energy
-
-
-def apply_causal_mask(
-    energy: Tensor, mask: Tensor, mask_value: Tensor, device: str
-) -> Tensor:
-    """Applies a causal mask to the energy tensor."""
-    i, j = energy.shape[-2:]
-    r = torch.arange(i, device=device)
-    mask = rearrange(r, "i -> () () i ()") < rearrange(r, "j -> () () () j")
-    mask = F.pad(mask, (j - i, 0), value=False)
-    energy.masked_fill_(mask, mask_value)
-    return energy
diff --git a/text_recognizer/networks/transformer/decoder.py b/text_recognizer/networks/transformer/decoder.py
deleted file mode 100644
index 826bc13..0000000
--- a/text_recognizer/networks/transformer/decoder.py
+++ /dev/null
@@ -1,41 +0,0 @@
-"""Transformer decoder module."""
-from copy import deepcopy
-from typing import Optional
-
-from torch import Tensor, nn
-
-from text_recognizer.networks.transformer.decoder_block import DecoderBlock
-from text_recognizer.networks.transformer.embeddings.rotary import RotaryEmbedding
-
-
-class Decoder(nn.Module):
-    """Decoder Network."""
-
-    def __init__(
-        self,
-        depth: int,
-        dim: int,
-        block: DecoderBlock,
-        rotary_embedding: RotaryEmbedding,
-    ) -> None:
-        super().__init__()
-        self.depth = depth
-        self.rotary_embedding = rotary_embedding
-        self.blocks = nn.ModuleList([deepcopy(block) for _ in range(self.depth)])
-        self.ln = nn.LayerNorm(dim)
-
-    def forward(
-        self,
-        x: Tensor,
-        context: Optional[Tensor] = None,
-        mask: Optional[Tensor] = None,
-    ) -> Tensor:
-        """Applies attention blocks."""
-        for block in self.blocks:
-            x = block(
-                x=x,
-                context=context,
-                mask=mask,
-                rotary_embedding=self.rotary_embedding,
-            )
-        return self.ln(x)
diff --git a/text_recognizer/networks/transformer/decoder_block.py b/text_recognizer/networks/transformer/decoder_block.py
deleted file mode 100644
index b8eb5c4..0000000
--- a/text_recognizer/networks/transformer/decoder_block.py
+++ /dev/null
@@ -1,44 +0,0 @@
-"""Transformer decoder module."""
-from copy import deepcopy
-from typing import Optional, Type
-
-from torch import Tensor, nn
-
-from text_recognizer.networks.transformer.attention import Attention
-from text_recognizer.networks.transformer.embeddings.rotary import RotaryEmbedding
-from text_recognizer.networks.transformer.ff import FeedForward
-
-
-class DecoderBlock(nn.Module):
-    """Residual decoder block."""
-
-    def __init__(
-        self,
-        self_attn: Attention,
-        norm: Type[nn.Module],
-        ff: FeedForward,
-        cross_attn: Optional[Attention] = None,
-    ) -> None:
-        super().__init__()
-        self.ln_attn = norm
-        self.attn = self_attn
-        self.ln_cross_attn = deepcopy(norm)
-        self.cross_attn = cross_attn
-        self.ln_ff = deepcopy(norm)
-        self.ff = ff
-
-    def forward(
-        self,
-        x: Tensor,
-        rotary_embedding: RotaryEmbedding,
-        context: Optional[Tensor] = None,
-        mask: Optional[Tensor] = None,
-    ) -> Tensor:
-        """Applies decoder block on input signals."""
-        x = x + self.attn(self.ln_attn(x), mask=mask, rotary_embedding=rotary_embedding)
-        x = x + self.cross_attn(
-            x=self.ln_cross_attn(x),
-            context=context,
-        )
-        x = x + self.ff(self.ln_ff(x))
-        return x
diff --git a/text_recognizer/networks/transformer/embeddings/__init__.py b/text_recognizer/networks/transformer/embeddings/__init__.py
deleted file mode 100644
index bb3f904..0000000
--- a/text_recognizer/networks/transformer/embeddings/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-"""Positional encodings for transformers."""
diff --git a/text_recognizer/networks/transformer/embeddings/axial.py b/text_recognizer/networks/transformer/embeddings/axial.py
deleted file mode 100644
index 9b872a9..0000000
--- a/text_recognizer/networks/transformer/embeddings/axial.py
+++ /dev/null
@@ -1,104 +0,0 @@
-"""Axial attention for multi-dimensional data.
-
-Stolen from:
-    https://github.com/lucidrains/axial-attention/blob/
-    eff2c10c2e76c735a70a6b995b571213adffbbb7/axial_attention/axial_attention.py#L100
-"""
-from functools import reduce
-from operator import mul
-from typing import Optional, Sequence
-
-import torch
-from torch import Tensor, nn
-
-
-class AxialPositionalEmbedding(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        axial_shape: Sequence[int],
-        axial_dims: Optional[Sequence[int]] = None,
-    ) -> None:
-        super().__init__()
-
-        self.dim = dim
-        self.shape = axial_shape
-        self.max_seq_len = reduce(mul, axial_shape, 1)
-
-        self.summed = axial_dims is None
-        axial_dims = ((dim,) * len(axial_shape)) if self.summed else axial_dims
-
-        assert len(self.shape) == len(
-            axial_dims
-        ), "number of axial dimensions must equal the number of dimensions in the shape"
-        assert (
-            self.summed or not self.summed and sum(axial_dims) == dim
-        ), f"axial dimensions must sum up to the target dimension {dim}"
-
-        self.weights = ParameterList(self, "weights", len(axial_shape))
-
-        for ind, (shape, axial_dim) in enumerate(zip(self.shape, axial_dims)):
-            ax_shape = [1] * len(self.shape)
-            ax_shape[ind] = shape
-            ax_shape = (1, *ax_shape, axial_dim)
-            ax_emb = nn.Parameter(torch.zeros(ax_shape).normal_(0, 1))
-            self.weights.append(ax_emb)
-
-    def forward(self, x: Tensor) -> Tensor:
-        """Returns axial positional embedding."""
-        b, t, _ = x.shape
-        assert (
-            t <= self.max_seq_len
-        ), f"Sequence length ({t}) must be less than the maximum sequence length allowed ({self.max_seq_len})"
-        embs = []
-
-        for ax_emb in self.weights.to_list():
-            axial_dim = ax_emb.shape[-1]
-            expand_shape = (b, *self.shape, axial_dim)
-            emb = ax_emb.expand(expand_shape).reshape(b, self.max_seq_len, axial_dim)
-            embs.append(emb)
-
-        pos_emb = sum(embs) if self.summed else torch.cat(embs, dim=-1)
-        return pos_emb[:, :t].to(x)
-
-
-# a mock parameter list object until below issue is resolved
-# https://github.com/pytorch/pytorch/issues/36035
-class ParameterList(object):
-    def __init__(self, kls, prefix, length):
-        self.ind = 0
-        self.kls = kls
-        self.prefix = prefix
-        self.length = length
-
-    def _keyname(self, prefix, ind):
-        return f"{prefix}_{ind}"
-
-    def append(self, x):
-        setattr(self.kls, self._keyname(self.prefix, self.ind), x)
-        self.ind += 1
-
-    def to_list(self):
-        return [
-            getattr(self.kls, self._keyname(self.prefix, i)) for i in range(self.length)
-        ]
-
-
-class AxialPositionalEmbeddingImage(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        axial_shape: Sequence[int],
-        axial_dims: Optional[Sequence[int]] = None,
-    ) -> None:
-        super().__init__()
-        axial_dims = (dim // 2, dim // 2) if axial_dims is None else axial_dims
-        assert len(axial_shape) == 2, "Axial shape must have 2 dimensions for images"
-        self.dim = dim
-        self.pos_emb = AxialPositionalEmbedding(dim, axial_shape, axial_dims)
-
-    def forward(self, img):
-        b, c, h, w = img.shape
-        img = img.permute(0, 2, 3, 1).reshape(b, h * w, c)
-        pos_emb = self.pos_emb(img)
-        return pos_emb.reshape(b, h, w, self.dim).permute(0, 3, 1, 2)
diff --git a/text_recognizer/networks/transformer/embeddings/rotary.py b/text_recognizer/networks/transformer/embeddings/rotary.py
deleted file mode 100644
index ca0a260..0000000
--- a/text_recognizer/networks/transformer/embeddings/rotary.py
+++ /dev/null
@@ -1,67 +0,0 @@
-"""Roatary embedding.
-
-Stolen from lucidrains:
-    https://github.com/lucidrains/rotary-embedding-torch
-
-Explanation of roatary:
-    https://blog.eleuther.ai/rotary-embeddings/
-"""
-from inspect import isfunction
-
-from einops import rearrange, repeat
-import torch
-from torch import Tensor, nn
-
-
-class RotaryEmbedding(nn.Module):
-    """Rotary positional embedding."""
-
-    def __init__(self, dim: int) -> None:
-        super().__init__()
-        inv_freqs = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
-        self.register_buffer("inv_freqs", inv_freqs)
-        self.cache = {}
-
-    def rotate(self, t: Tensor, dim: int = -2) -> Tensor:
-        """Rotate vector."""
-        device, n = t.device, t.shape[dim]
-        freqs = self.forward(lambda: torch.arange(n, device=device), cache_key=n)
-        return apply_rotary_emb(t, freqs)
-
-    def forward(self, t: Tensor, cache_key: int) -> Tensor:
-        """Encodes tensor x with rotary embeddings."""
-        if cache_key in self.cache:
-            return self.cache[cache_key]
-
-        if isfunction(t):
-            t = t()
-
-        freqs = self.inv_freqs
-        freqs = torch.einsum("..., f -> ... f", t.type(freqs.dtype), freqs)
-        freqs = repeat(freqs, "... n -> ... (n r)", r=2)
-        self.cache[cache_key] = freqs
-        return freqs
-
-
-def rotate_half(x: Tensor) -> Tensor:
-    x = rearrange(x, "... (d r) -> ... d r", r=2)
-    x1, x2 = x.unbind(dim=-1)
-    x = torch.stack((-x2, x1), dim=-1)
-    return rearrange(x, "... d r -> ... (d r)")
-
-
-def apply_rotary_emb(t: Tensor, freqs: Tensor, start_index: int = 0) -> Tensor:
-    freqs = freqs.to(t)
-    rot_dim = freqs.shape[-1]
-    end_index = start_index + rot_dim
-    assert rot_dim <= t.shape[-1], (
-        f"feature dimension {t.shape[-1]} is not of sufficient size to rotate"
-        f"in all the positions {rot_dim}"
-    )
-    t_left, t, t_right = (
-        t[..., :start_index],
-        t[..., start_index:end_index],
-        t[..., end_index:],
-    )
-    t = (t * freqs.cos()) + (rotate_half(t) * freqs.sin())
-    return torch.cat((t_left, t, t_right), dim=-1)
diff --git a/text_recognizer/networks/transformer/ff.py b/text_recognizer/networks/transformer/ff.py
deleted file mode 100644
index 3ccf5b5..0000000
--- a/text_recognizer/networks/transformer/ff.py
+++ /dev/null
@@ -1,45 +0,0 @@
-"""Feedforward layer in transformer.
-
-Stolen from lucidrains:
-    https://github.com/lucidrains/x-transformers/blob/main/x_transformers/x_transformers.py
-"""
-from typing import Optional
-
-import torch.nn.functional as F
-from torch import Tensor, nn
-
-
-class GEGLU(nn.Module):
-    def __init__(self, dim_in: int, dim_out: int) -> None:
-        super().__init__()
-        self.fc = nn.Linear(dim_in, dim_out * 2)
-
-    def forward(self, x: Tensor) -> Tensor:
-        x, gate = self.fc(x).chunk(2, dim=-1)
-        return x * F.gelu(gate)
-
-
-class FeedForward(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        dim_out: Optional[int] = None,
-        expansion_factor: int = 4,
-        glu: bool = True,
-        dropout_rate: float = 0.0,
-    ) -> None:
-        super().__init__()
-        inner_dim = dim * expansion_factor
-        dim_out = dim_out if dim_out is not None else dim
-        in_projection = (
-            nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU())
-            if not glu
-            else GEGLU(dim, inner_dim)
-        )
-
-        self.mlp = nn.Sequential(
-            in_projection, nn.Dropout(dropout_rate), nn.Linear(inner_dim, dim_out)
-        )
-
-    def forward(self, x: Tensor) -> Tensor:
-        return self.mlp(x)
diff --git a/text_recognizer/networks/transformer/norm.py b/text_recognizer/networks/transformer/norm.py
deleted file mode 100644
index 1431327..0000000
--- a/text_recognizer/networks/transformer/norm.py
+++ /dev/null
@@ -1,51 +0,0 @@
-"""Normalization layers for transfromers.
-
-Copied from lucidrains:
-    https://github.com/lucidrains/x-transformers/blob/main/x_transformers/x_transformers.py
-
-"""
-from typing import Optional, Type
-
-import torch
-from torch import Tensor, nn
-
-
-class RMSNorm(nn.Module):
-    """Root mean square layer normalization."""
-
-    def __init__(self, dim: int, eps: float = 1e-8) -> None:
-        super().__init__()
-        self.scale = dim**-0.5
-        self.eps = eps
-        self.g = nn.Parameter(torch.ones(dim))
-
-    def forward(self, x: Tensor) -> Tensor:
-        """Applies normalization."""
-        norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
-        return x / norm.clamp(min=self.eps) * self.g
-
-
-class PreNorm(nn.Module):
-    """Applies layer normalization then function."""
-
-    def __init__(
-        self,
-        normalized_shape: int,
-        fn: Type[nn.Module],
-        context_dim: Optional[int] = None,
-    ) -> None:
-        super().__init__()
-        self.norm = nn.LayerNorm(normalized_shape)
-        self.fn = fn
-        self.norm_context = (
-            nn.LayerNorm(context_dim) if context_dim is not None else None
-        )
-
-    def forward(self, x: Tensor, **kwargs) -> Tensor:
-        """Applies pre norm."""
-        x = self.norm(x)
-        if self.norm_context is not None:
-            context = kwargs["context"]
-            normed_context = self.norm_context(context)
-            kwargs.update(context=normed_context)
-        return self.fn(x, **kwargs)
diff --git a/text_recognizer/optimizer/__init__.py b/text_recognizer/optimizer/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/text_recognizer/optimizer/laprop.py b/text_recognizer/optimizer/laprop.py
new file mode 100644
index 0000000..9b4f4d4
--- /dev/null
+++ b/text_recognizer/optimizer/laprop.py
@@ -0,0 +1,129 @@
+"""https://github.com/Z-T-WANG/LaProp-Optimizer/blob/master/laprop.py"""
+from torch.optim import Optimizer
+import math
+import torch
+
+
+class LaProp(Optimizer):
+    def __init__(
+        self,
+        params,
+        lr=4e-4,
+        betas=(0.9, 0.999),
+        eps=1e-15,
+        weight_decay=0,
+        amsgrad=False,
+        centered=False,
+    ):
+
+        self.steps_before_using_centered = 10
+
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            amsgrad=amsgrad,
+            centered=centered,
+        )
+        super(LaProp, self).__init__(params, defaults)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError(
+                        "Adam does not support sparse gradients, please consider SparseAdam instead"
+                    )
+                amsgrad = group["amsgrad"]
+                centered = group["centered"]
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state["step"] = 0
+                    # Exponential moving average of gradient values
+                    state["exp_avg"] = torch.zeros_like(p.data)
+                    # Exponential moving average of learning rates
+                    state["exp_avg_lr_1"] = 0.0
+                    state["exp_avg_lr_2"] = 0.0
+                    # Exponential moving average of squared gradient values
+                    state["exp_avg_sq"] = torch.zeros_like(p.data)
+                    if centered:
+                        # Exponential moving average of gradient values as calculated by beta2
+                        state["exp_mean_avg_beta2"] = torch.zeros_like(p.data)
+                    if amsgrad:
+                        # Maintains max of all exp. moving avg. of sq. grad. values
+                        state["max_exp_avg_sq"] = torch.zeros_like(p.data)
+
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+                if centered:
+                    exp_mean_avg_beta2 = state["exp_mean_avg_beta2"]
+                if amsgrad:
+                    max_exp_avg_sq = state["max_exp_avg_sq"]
+                beta1, beta2 = group["betas"]
+
+                state["step"] += 1
+
+                # Decay the first and second moment running average coefficient
+                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
+
+                state["exp_avg_lr_1"] = (
+                    state["exp_avg_lr_1"] * beta1 + (1 - beta1) * group["lr"]
+                )
+                state["exp_avg_lr_2"] = state["exp_avg_lr_2"] * beta2 + (1 - beta2)
+
+                bias_correction1 = (
+                    state["exp_avg_lr_1"] / group["lr"] if group["lr"] != 0.0 else 1.0
+                )  # 1 - beta1 ** state['step']
+                step_size = 1 / bias_correction1
+
+                bias_correction2 = state["exp_avg_lr_2"]
+
+                denom = exp_avg_sq
+                if centered:
+                    exp_mean_avg_beta2.mul_(beta2).add_(1 - beta2, grad)
+                    if state["step"] > self.steps_before_using_centered:
+                        mean = exp_mean_avg_beta2**2
+                        denom = denom - mean
+
+                if amsgrad:
+                    if not (
+                        centered and state["step"] <= self.steps_before_using_centered
+                    ):
+                        # Maintains the maximum of all (centered) 2nd moment running avg. till now
+                        torch.max(max_exp_avg_sq, denom, out=max_exp_avg_sq)
+                        # Use the max. for normalizing running avg. of gradient
+                        denom = max_exp_avg_sq
+
+                denom = denom.div(bias_correction2).sqrt_().add_(group["eps"])
+                step_of_this_grad = grad / denom
+                exp_avg.mul_(beta1).add_((1 - beta1) * group["lr"], step_of_this_grad)
+
+                p.data.add_(-step_size, exp_avg)
+                if group["weight_decay"] != 0:
+                    p.data.add_(-group["weight_decay"], p.data)
+
+        return loss
diff --git a/text_recognizer/optimizers/__init__.py b/text_recognizer/optimizers/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/text_recognizer/optimizers/laprop.py b/text_recognizer/optimizers/laprop.py
deleted file mode 100644
index 9b4f4d4..0000000
--- a/text_recognizer/optimizers/laprop.py
+++ /dev/null
@@ -1,129 +0,0 @@
-"""https://github.com/Z-T-WANG/LaProp-Optimizer/blob/master/laprop.py"""
-from torch.optim import Optimizer
-import math
-import torch
-
-
-class LaProp(Optimizer):
-    def __init__(
-        self,
-        params,
-        lr=4e-4,
-        betas=(0.9, 0.999),
-        eps=1e-15,
-        weight_decay=0,
-        amsgrad=False,
-        centered=False,
-    ):
-
-        self.steps_before_using_centered = 10
-
-        if not 0.0 <= lr:
-            raise ValueError("Invalid learning rate: {}".format(lr))
-        if not 0.0 <= eps:
-            raise ValueError("Invalid epsilon value: {}".format(eps))
-        if not 0.0 <= betas[0] < 1.0:
-            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
-        if not 0.0 <= betas[1] < 1.0:
-            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
-        defaults = dict(
-            lr=lr,
-            betas=betas,
-            eps=eps,
-            weight_decay=weight_decay,
-            amsgrad=amsgrad,
-            centered=centered,
-        )
-        super(LaProp, self).__init__(params, defaults)
-
-    def step(self, closure=None):
-        """Performs a single optimization step.
-
-        Arguments:
-            closure (callable, optional): A closure that reevaluates the model
-                and returns the loss.
-        """
-        loss = None
-        if closure is not None:
-            loss = closure()
-
-        for group in self.param_groups:
-            for p in group["params"]:
-                if p.grad is None:
-                    continue
-                grad = p.grad.data
-                if grad.is_sparse:
-                    raise RuntimeError(
-                        "Adam does not support sparse gradients, please consider SparseAdam instead"
-                    )
-                amsgrad = group["amsgrad"]
-                centered = group["centered"]
-
-                state = self.state[p]
-
-                # State initialization
-                if len(state) == 0:
-                    state["step"] = 0
-                    # Exponential moving average of gradient values
-                    state["exp_avg"] = torch.zeros_like(p.data)
-                    # Exponential moving average of learning rates
-                    state["exp_avg_lr_1"] = 0.0
-                    state["exp_avg_lr_2"] = 0.0
-                    # Exponential moving average of squared gradient values
-                    state["exp_avg_sq"] = torch.zeros_like(p.data)
-                    if centered:
-                        # Exponential moving average of gradient values as calculated by beta2
-                        state["exp_mean_avg_beta2"] = torch.zeros_like(p.data)
-                    if amsgrad:
-                        # Maintains max of all exp. moving avg. of sq. grad. values
-                        state["max_exp_avg_sq"] = torch.zeros_like(p.data)
-
-                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
-                if centered:
-                    exp_mean_avg_beta2 = state["exp_mean_avg_beta2"]
-                if amsgrad:
-                    max_exp_avg_sq = state["max_exp_avg_sq"]
-                beta1, beta2 = group["betas"]
-
-                state["step"] += 1
-
-                # Decay the first and second moment running average coefficient
-                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
-
-                state["exp_avg_lr_1"] = (
-                    state["exp_avg_lr_1"] * beta1 + (1 - beta1) * group["lr"]
-                )
-                state["exp_avg_lr_2"] = state["exp_avg_lr_2"] * beta2 + (1 - beta2)
-
-                bias_correction1 = (
-                    state["exp_avg_lr_1"] / group["lr"] if group["lr"] != 0.0 else 1.0
-                )  # 1 - beta1 ** state['step']
-                step_size = 1 / bias_correction1
-
-                bias_correction2 = state["exp_avg_lr_2"]
-
-                denom = exp_avg_sq
-                if centered:
-                    exp_mean_avg_beta2.mul_(beta2).add_(1 - beta2, grad)
-                    if state["step"] > self.steps_before_using_centered:
-                        mean = exp_mean_avg_beta2**2
-                        denom = denom - mean
-
-                if amsgrad:
-                    if not (
-                        centered and state["step"] <= self.steps_before_using_centered
-                    ):
-                        # Maintains the maximum of all (centered) 2nd moment running avg. till now
-                        torch.max(max_exp_avg_sq, denom, out=max_exp_avg_sq)
-                        # Use the max. for normalizing running avg. of gradient
-                        denom = max_exp_avg_sq
-
-                denom = denom.div(bias_correction2).sqrt_().add_(group["eps"])
-                step_of_this_grad = grad / denom
-                exp_avg.mul_(beta1).add_((1 - beta1) * group["lr"], step_of_this_grad)
-
-                p.data.add_(-step_size, exp_avg)
-                if group["weight_decay"] != 0:
-                    p.data.add_(-group["weight_decay"], p.data)
-
-        return loss
-- 
cgit v1.2.3-70-g09d2