diff options
Diffstat (limited to 'text_recognizer')
-rw-r--r-- | text_recognizer/criterions/barlow_twins.py | 26 | ||||
-rw-r--r-- | text_recognizer/data/transforms/barlow.py | 19 | ||||
-rw-r--r-- | text_recognizer/models/barlow_twins.py | 45 | ||||
-rw-r--r-- | text_recognizer/networks/barlow_twins/__init__.py | 1 | ||||
-rw-r--r-- | text_recognizer/networks/barlow_twins/network.py | 18 | ||||
-rw-r--r-- | text_recognizer/networks/barlow_twins/projector.py | 36 |
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) |