Remove Barlow Twins

6 files changed, 0 insertions, 145 deletions

diff --git a/text_recognizer/criterions/barlow_twins.py b/text_recognizer/criterions/barlow_twins.py
deleted file mode 100644
index fe30b22..0000000
--- a/text_recognizer/criterions/barlow_twins.py
+++ /dev/null
@@ -1,26 +0,0 @@
-"""Barlow twins loss function."""
-
-import torch
-from torch import nn, Tensor
-
-
-def off_diagonal(x: Tensor) -> Tensor:
-    n, m = x.shape
-    assert n == m
-    return x.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten()
-
-
-class BarlowTwinsLoss(nn.Module):
-    def __init__(self, dim: int, lambda_: float) -> None:
-        super().__init__()
-        self.bn = nn.BatchNorm1d(dim, affine=False)
-        self.lambda_ = lambda_
-
-    def forward(self, z1: Tensor, z2: Tensor) -> Tensor:
-        """Calculates the Barlow Twin loss."""
-        c = self.bn(z1).T @ self.bn(z2)
-        c.div_(z1.shape[0])
-
-        on_diag = torch.diagonal(c).add_(-1).pow_(2).sum()
-        off_diag = off_diagonal(c).pow_(2).sum()
-        return on_diag + self.lambda_ * off_diag
diff --git a/text_recognizer/data/transforms/barlow.py b/text_recognizer/data/transforms/barlow.py
deleted file mode 100644
index 78683cb..0000000
--- a/text_recognizer/data/transforms/barlow.py
+++ /dev/null
@@ -1,19 +0,0 @@
-"""Augmentations for training Barlow Twins."""
-from omegaconf.dictconfig import DictConfig
-from torch import Tensor
-
-from text_recognizer.data.transforms.load_transform import load_transform
-
-
-class BarlowTransform:
-    """Applies two different transforms to input data."""
-
-    def __init__(self, prim: DictConfig, bis: DictConfig) -> None:
-        self.prim = load_transform(prim)
-        self.bis = load_transform(bis)
-
-    def __call__(self, data: Tensor) -> Tensor:
-        """Applies two different augmentation on the input."""
-        x_prim = self.prim(data)
-        x_bis = self.bis(data)
-        return x_prim, x_bis
diff --git a/text_recognizer/models/barlow_twins.py b/text_recognizer/models/barlow_twins.py
deleted file mode 100644
index 6e2719d..0000000
--- a/text_recognizer/models/barlow_twins.py
+++ /dev/null
@@ -1,45 +0,0 @@
-"""PyTorch Lightning Barlow Twins model."""
-from typing import Tuple, Type
-import attr
-from torch import nn
-from torch import Tensor
-
-from text_recognizer.models.base import BaseLitModel
-from text_recognizer.criterions.barlow_twins import BarlowTwinsLoss
-
-
-@attr.s(auto_attribs=True, eq=False)
-class BarlowTwinsLitModel(BaseLitModel):
-    """Barlow Twins training proceduer."""
-
-    network: Type[nn.Module] = attr.ib()
-    loss_fn: BarlowTwinsLoss = attr.ib()
-
-    def forward(self, data: Tensor) -> Tensor:
-        """Encodes image to projector latent."""
-        return self.network(data)
-
-    def training_step(self, batch: Tuple[Tensor, Tensor], batch_idx: int) -> Tensor:
-        """Training step."""
-        data, _ = batch
-        x1, x2 = data
-        z1, z2 = self(x1), self(x2)
-        loss = self.loss_fn(z1, z2)
-        self.log("train/loss", loss)
-        return loss
-
-    def validation_step(self, batch: Tuple[Tensor, Tensor], batch_idx: int) -> None:
-        """Validation step."""
-        data, _ = batch
-        x1, x2 = data
-        z1, z2 = self(x1), self(x2)
-        loss = self.loss_fn(z1, z2)
-        self.log("val/loss", loss, prog_bar=True)
-
-    def test_step(self, batch: Tuple[Tensor, Tensor], batch_idx: int) -> None:
-        """Test step."""
-        data, _ = batch
-        x1, x2 = data
-        z1, z2 = self(x1), self(x2)
-        loss = self.loss_fn(z1, z2)
-        self.log("test/loss", loss, prog_bar=True)
diff --git a/text_recognizer/networks/barlow_twins/__init__.py b/text_recognizer/networks/barlow_twins/__init__.py
deleted file mode 100644
index 0b74818..0000000
--- a/text_recognizer/networks/barlow_twins/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-"""Module for projector network in Barlow Twins."""
diff --git a/text_recognizer/networks/barlow_twins/network.py b/text_recognizer/networks/barlow_twins/network.py
deleted file mode 100644
index a3e3750..0000000
--- a/text_recognizer/networks/barlow_twins/network.py
+++ /dev/null
@@ -1,18 +0,0 @@
-"""Barlow Twins network."""
-from typing import Type
-
-from torch import nn, Tensor
-import torch.nn.functional as F
-
-
-class BarlowTwins(nn.Module):
-    def __init__(self, encoder: Type[nn.Module], projector: Type[nn.Module]) -> None:
-        super().__init__()
-        self.encoder = encoder
-        self.projector = projector
-
-    def forward(self, x: Tensor) -> Tensor:
-        z = self.encoder(x)
-        z_e = F.adaptive_avg_pool2d(z, (1, 1)).flatten(start_dim=1)
-        z_p = self.projector(z_e)
-        return z_p
diff --git a/text_recognizer/networks/barlow_twins/projector.py b/text_recognizer/networks/barlow_twins/projector.py
deleted file mode 100644
index 05d5e2e..0000000
--- a/text_recognizer/networks/barlow_twins/projector.py
+++ /dev/null
@@ -1,36 +0,0 @@
-"""Projector network in Barlow Twins."""
-
-from typing import List
-import torch
-from torch import nn
-from torch import Tensor
-
-
-class Projector(nn.Module):
-    """MLP network."""
-
-    def __init__(self, dims: List[int]) -> None:
-        super().__init__()
-        self.dims = dims
-        self.network = self._build()
-
-    def _build(self) -> nn.Sequential:
-        """Builds projector network."""
-        layers = [
-            nn.Sequential(
-                nn.Linear(
-                    in_features=self.dims[i], out_features=self.dims[i + 1], bias=False
-                ),
-                nn.BatchNorm1d(self.dims[i + 1]),
-                nn.ReLU(inplace=True),
-            )
-            for i in range(len(self.dims) - 2)
-        ]
-        layers.append(
-            nn.Linear(in_features=self.dims[-2], out_features=self.dims[-1], bias=False)
-        )
-        return nn.Sequential(*layers)
-
-    def forward(self, x: Tensor) -> Tensor:
-        """Project latent to higher dimesion."""
-        return self.network(x)