From 9a8044f4a3826a119416665741b709cd686fca87 Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Wed, 27 Oct 2021 22:27:24 +0200 Subject: Remove Barlow Twins --- text_recognizer/criterions/barlow_twins.py | 26 ------ text_recognizer/data/transforms/barlow.py | 19 ---- text_recognizer/models/barlow_twins.py | 45 --------- text_recognizer/networks/barlow_twins/__init__.py | 1 - text_recognizer/networks/barlow_twins/network.py | 18 ---- text_recognizer/networks/barlow_twins/projector.py | 36 ------- .../datamodule/transform/barlow_paragraphs.yaml | 46 --------- training/conf/experiment/barlow_twins.yaml | 102 -------------------- .../conf/experiment/barlow_twins_paragraphs.yaml | 103 --------------------- 9 files changed, 396 deletions(-) delete mode 100644 text_recognizer/criterions/barlow_twins.py delete mode 100644 text_recognizer/data/transforms/barlow.py delete mode 100644 text_recognizer/models/barlow_twins.py delete mode 100644 text_recognizer/networks/barlow_twins/__init__.py delete mode 100644 text_recognizer/networks/barlow_twins/network.py delete mode 100644 text_recognizer/networks/barlow_twins/projector.py delete mode 100644 training/conf/datamodule/transform/barlow_paragraphs.yaml delete mode 100644 training/conf/experiment/barlow_twins.yaml delete mode 100644 training/conf/experiment/barlow_twins_paragraphs.yaml 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 -- cgit v1.2.3-70-g09d2