summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
Diffstat
-rw-r--r--text_recognizer/criterions/barlow_twins.py26
-rw-r--r--text_recognizer/data/transforms/barlow.py19
-rw-r--r--text_recognizer/models/barlow_twins.py45
-rw-r--r--text_recognizer/networks/barlow_twins/__init__.py1
-rw-r--r--text_recognizer/networks/barlow_twins/network.py18
-rw-r--r--text_recognizer/networks/barlow_twins/projector.py36
-rw-r--r--training/conf/datamodule/transform/barlow_paragraphs.yaml46
-rw-r--r--training/conf/experiment/barlow_twins.yaml102
-rw-r--r--training/conf/experiment/barlow_twins_paragraphs.yaml103
9 files changed, 0 insertions, 396 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)
diff --git a/training/conf/datamodule/transform/barlow_paragraphs.yaml b/training/conf/datamodule/transform/barlow_paragraphs.yaml
deleted file mode 100644
index 5eefce5..0000000
--- a/training/conf/datamodule/transform/barlow_paragraphs.yaml
+++ /dev/null
@@ -1,46 +0,0 @@
-
-barlow:
- _target_: text_recognizer.data.transforms.barlow.BarlowTransform
- prim:
- random_crop:
- _target_: torchvision.transforms.RandomCrop
- size: [576, 640]
- padding: null
- pad_if_needed: true
- fill: 0
- padding_mode: constant
-
- color_jitter:
- _target_: torchvision.transforms.ColorJitter
- brightness: [0.8, 1.6]
-
- random_affine:
- _target_: torchvision.transforms.RandomAffine
- degrees: 1
- shear: [-10, 10]
- interpolation: BILINEAR
-
- to_tensor:
- _target_: torchvision.transforms.ToTensor
-
- bis:
- random_crop:
- _target_: torchvision.transforms.RandomCrop
- size: [576, 640]
- padding: null
- pad_if_needed: true
- fill: 0
- padding_mode: constant
-
- color_jitter:
- _target_: torchvision.transforms.ColorJitter
- brightness: [0.8, 2.0]
-
- random_affine:
- _target_: torchvision.transforms.RandomAffine
- degrees: 1
- shear: [-5, 5]
- interpolation: BILINEAR
-
- to_tensor:
- _target_: torchvision.transforms.ToTensor
diff --git a/training/conf/experiment/barlow_twins.yaml b/training/conf/experiment/barlow_twins.yaml
deleted file mode 100644
index cc1295d..0000000
--- a/training/conf/experiment/barlow_twins.yaml
+++ /dev/null
@@ -1,102 +0,0 @@
-# @package _global_
-
-defaults:
- - override /criterion: null
- - override /datamodule: null
- - override /network: null
- - override /model: null
- - override /lr_schedulers: null
- - override /optimizers: null
-
-epochs: &epochs 1000
-summary: [[1, 1, 56, 1024]]
-
-criterion:
- _target_: text_recognizer.criterions.barlow_twins.BarlowTwinsLoss
- dim: 512
- lambda_: 3.9e-3
-
-callbacks:
- stochastic_weight_averaging:
- _target_: pytorch_lightning.callbacks.StochasticWeightAveraging
- swa_epoch_start: 0.75
- swa_lrs: 1.0e-5
- annealing_epochs: 10
- annealing_strategy: cos
- device: null
-
-optimizers:
- madgrad:
- _target_: madgrad.MADGRAD
- lr: 1.0e-3
- momentum: 0.9
- weight_decay: 1.0e-6
- eps: 1.0e-6
- parameters: network
-
-lr_schedulers:
- network:
- _target_: torch.optim.lr_scheduler.OneCycleLR
- max_lr: 3.0e-4
- total_steps: null
- epochs: *epochs
- steps_per_epoch: 45
- pct_start: 0.03
- anneal_strategy: cos
- cycle_momentum: true
- base_momentum: 0.85
- max_momentum: 0.95
- div_factor: 25
- final_div_factor: 1.0e4
- three_phase: false
- last_epoch: -1
- verbose: false
- # Non-class arguments
- interval: step
- monitor: val/loss
-
-datamodule:
- _target_: text_recognizer.data.iam_lines.IAMLines
- batch_size: 16
- num_workers: 12
- train_fraction: 0.9
- pin_memory: false
- transform: transform/iam_lines_barlow.yaml
- test_transform: transform/iam_lines_barlow.yaml
- mapping:
- _target_: text_recognizer.data.mappings.emnist_mapping.EmnistMapping
-
-network:
- _target_: text_recognizer.networks.barlow_twins.network.BarlowTwins
- encoder:
- _target_: text_recognizer.networks.encoders.efficientnet.EfficientNet
- arch: b0
- out_channels: 1280
- stochastic_dropout_rate: 0.2
- bn_momentum: 0.99
- bn_eps: 1.0e-3
- projector:
- _target_: text_recognizer.networks.barlow_twins.projector.Projector
- dims: [1280, 512, 512, 512]
-
-model:
- _target_: text_recognizer.models.barlow_twins.BarlowTwinsLitModel
-
-trainer:
- _target_: pytorch_lightning.Trainer
- stochastic_weight_avg: true
- auto_scale_batch_size: binsearch
- auto_lr_find: false
- gradient_clip_val: 0.0
- fast_dev_run: false
- gpus: 1
- precision: 16
- max_epochs: *epochs
- terminate_on_nan: true
- weights_summary: null
- limit_train_batches: 1.0
- limit_val_batches: 1.0
- limit_test_batches: 1.0
- resume_from_checkpoint: null
- accumulate_grad_batches: 32
- overfit_batches: 0
diff --git a/training/conf/experiment/barlow_twins_paragraphs.yaml b/training/conf/experiment/barlow_twins_paragraphs.yaml
deleted file mode 100644
index 9552c0b..0000000
--- a/training/conf/experiment/barlow_twins_paragraphs.yaml
+++ /dev/null
@@ -1,103 +0,0 @@
-# @package _global_
-
-defaults:
- - override /criterion: null
- - override /datamodule: null
- - override /network: null
- - override /model: null
- - override /lr_schedulers: null
- - override /optimizers: null
-
-epochs: &epochs 1000
-summary: [[1, 1, 576, 640]]
-
-criterion:
- _target_: text_recognizer.criterions.barlow_twins.BarlowTwinsLoss
- dim: 512
- lambda_: 3.9e-3
-
-# callbacks:
-# stochastic_weight_averaging:
-# _target_: pytorch_lightning.callbacks.StochasticWeightAveraging
-# swa_epoch_start: 0.75
-# swa_lrs: 1.0e-5
-# annealing_epochs: 10
-# annealing_strategy: cos
-# device: null
-
-optimizers:
- madgrad:
- _target_: madgrad.MADGRAD
- lr: 1.0e-3
- momentum: 0.9
- weight_decay: 1.0e-6
- eps: 1.0e-6
- parameters: network
-
-lr_schedulers:
- network:
- _target_: torch.optim.lr_scheduler.OneCycleLR
- max_lr: 1.0e-3
- total_steps: null
- epochs: *epochs
- steps_per_epoch: 40
- pct_start: 0.03
- anneal_strategy: cos
- cycle_momentum: true
- base_momentum: 0.85
- max_momentum: 0.95
- div_factor: 25
- final_div_factor: 1.0e4
- three_phase: false
- last_epoch: -1
- verbose: false
- # Non-class arguments
- interval: step
- monitor: val/loss
-
-datamodule:
- _target_: text_recognizer.data.iam_extended_paragraphs.IAMExtendedParagraphs
- batch_size: 4
- num_workers: 12
- train_fraction: 0.9
- pin_memory: true
- transform: transform/barlow_paragraphs.yaml
- test_transform: transform/barlow_paragraphs.yaml
- mapping:
- _target_: text_recognizer.data.mappings.emnist_mapping.EmnistMapping
- extra_symbols: [ "\n" ]
-
-network:
- _target_: text_recognizer.networks.barlow_twins.network.BarlowTwins
- encoder:
- _target_: text_recognizer.networks.encoders.efficientnet.EfficientNet
- arch: b0
- out_channels: 1280
- stochastic_dropout_rate: 0.2
- bn_momentum: 0.99
- bn_eps: 1.0e-3
- projector:
- _target_: text_recognizer.networks.barlow_twins.projector.Projector
- dims: [1280, 512, 512, 512]
-
-model:
- _target_: text_recognizer.models.barlow_twins.BarlowTwinsLitModel
-
-trainer:
- _target_: pytorch_lightning.Trainer
- stochastic_weight_avg: true
- auto_scale_batch_size: binsearch
- auto_lr_find: false
- gradient_clip_val: 0.0
- fast_dev_run: false
- gpus: 1
- precision: 16
- max_epochs: *epochs
- terminate_on_nan: true
- weights_summary: null
- limit_train_batches: 1.0
- limit_val_batches: 1.0
- limit_test_batches: 1.0
- resume_from_checkpoint: null
- accumulate_grad_batches: 128
- overfit_batches: 0
generated by cgit v1.2.3-70-g09d2 (git 2.46.0) at 2024-12-25 05:27:17 +0000