chore: remove axial attention

chore: remove axial attention

1 files changed, 0 insertions, 80 deletions

diff --git a/text_recognizer/networks/transformer/axial_attention/encoder.py b/text_recognizer/networks/transformer/axial_attention/encoder.py
deleted file mode 100644
index 3082bd6..0000000
--- a/text_recognizer/networks/transformer/axial_attention/encoder.py
+++ /dev/null
@@ -1,80 +0,0 @@
-"""Axial transformer encoder."""
-
-from typing import List
-
-import attr
-from torch import nn, Tensor
-
-from text_recognizer.networks.transformer.axial_attention.self_attention import (
-    SelfAttention,
-)
-from text_recognizer.networks.transformer.axial_attention.utils import (
-    calculate_permutations,
-    PermuteToForm,
-    Sequential,
-)
-from text_recognizer.networks.transformer.norm import PreNorm
-
-
-@attr.s(eq=False)
-class AxialEncoder(nn.Module):
-    """Axial transfomer encoder."""
-
-    def __attrs_pre_init__(self) -> None:
-        super().__init__()
-
-    shape: List[int] = attr.ib()
-    dim: int = attr.ib()
-    depth: int = attr.ib()
-    heads: int = attr.ib()
-    dim_head: int = attr.ib()
-    dim_index: int = attr.ib()
-    fn: nn.Sequential = attr.ib(init=False)
-
-    def __attrs_post_init__(self) -> None:
-        self._build()
-
-    def _build(self) -> None:
-        permutations = calculate_permutations(2, self.dim_index)
-        get_ff = lambda: nn.Sequential(
-            nn.LayerNorm([self.dim, *self.shape]),
-            nn.Conv2d(
-                in_channels=self.dim,
-                out_channels=4 * self.dim,
-                kernel_size=3,
-                padding=1,
-            ),
-            nn.Mish(inplace=True),
-            nn.Conv2d(
-                in_channels=4 * self.dim,
-                out_channels=self.dim,
-                kernel_size=3,
-                padding=1,
-            ),
-        )
-
-        layers = nn.ModuleList([])
-        for _ in range(self.depth):
-            attns = nn.ModuleList(
-                [
-                    PermuteToForm(
-                        permutation=permutation,
-                        fn=PreNorm(
-                            self.dim,
-                            SelfAttention(
-                                dim=self.dim, heads=self.heads, dim_head=self.dim_head
-                            ),
-                        ),
-                    )
-                    for permutation in permutations
-                ]
-            )
-            convs = nn.ModuleList([get_ff(), get_ff()])
-            layers.append(attns)
-            layers.append(convs)
-
-        self.fn = Sequential(layers)
-
-    def forward(self, x: Tensor) -> Tensor:
-        """Applies fn to input."""
-        return self.fn(x)