summaryrefslogtreecommitdiff
path: root/text_recognizer
diff options
context:
space:
mode:
Diffstat (limited to 'text_recognizer')
-rw-r--r--text_recognizer/models/base.py20
-rw-r--r--text_recognizer/networks/loss/__init__.py2
-rw-r--r--text_recognizer/networks/loss/loss.py32
-rw-r--r--text_recognizer/networks/util.py27
-rw-r--r--text_recognizer/networks/vq_transformer.py150
5 files changed, 13 insertions, 218 deletions
diff --git a/text_recognizer/models/base.py b/text_recognizer/models/base.py
index 46e5136..2d6e435 100644
--- a/text_recognizer/models/base.py
+++ b/text_recognizer/models/base.py
@@ -1,5 +1,5 @@
"""Base PyTorch Lightning model."""
-from typing import Any, Dict, Tuple, Type
+from typing import Any, Dict, List, Tuple, Type
import madgrad
import pytorch_lightning as pl
@@ -40,7 +40,7 @@ class LitBaseModel(pl.LightningModule):
args = {} or criterion_args["args"]
return getattr(nn, criterion_args["type"])(**args)
- def configure_optimizer(self) -> Dict[str, Any]:
+ def configure_optimizer(self) -> Tuple[List[type], List[Dict[str, Any]]]:
"""Configures optimizer and lr scheduler."""
args = {} or self.optimizer_args["args"]
if self.optimizer_args["type"] == "MADGRAD":
@@ -48,15 +48,15 @@ class LitBaseModel(pl.LightningModule):
else:
optimizer = getattr(torch.optim, self.optimizer_args["type"])(**args)
+ scheduler = {"monitor": self.monitor}
args = {} or self.lr_scheduler_args["args"]
- scheduler = getattr(torch.optim.lr_scheduler, self.lr_scheduler_args["type"])(
- **args
- )
- return {
- "optimizer": optimizer,
- "lr_scheduler": scheduler,
- "monitor": self.monitor,
- }
+ if "interval" in args:
+ scheduler["interval"] = args.pop("interval")
+
+ scheduler["scheduler"] = getattr(
+ torch.optim.lr_scheduler, self.lr_scheduler_args["type"]
+ )(**args)
+ return [optimizer], [scheduler]
def forward(self, data: Tensor) -> Tensor:
"""Feedforward pass."""
diff --git a/text_recognizer/networks/loss/__init__.py b/text_recognizer/networks/loss/__init__.py
index b489264..cb83608 100644
--- a/text_recognizer/networks/loss/__init__.py
+++ b/text_recognizer/networks/loss/__init__.py
@@ -1,2 +1,2 @@
"""Loss module."""
-from .loss import EmbeddingLoss, LabelSmoothingCrossEntropy
+from .loss import LabelSmoothingCrossEntropy
diff --git a/text_recognizer/networks/loss/loss.py b/text_recognizer/networks/loss/loss.py
index cf9fa0d..d12dc9c 100644
--- a/text_recognizer/networks/loss/loss.py
+++ b/text_recognizer/networks/loss/loss.py
@@ -1,39 +1,9 @@
"""Implementations of custom loss functions."""
-from pytorch_metric_learning import distances, losses, miners, reducers
import torch
from torch import nn
from torch import Tensor
-from torch.autograd import Variable
-import torch.nn.functional as F
-__all__ = ["EmbeddingLoss", "LabelSmoothingCrossEntropy"]
-
-
-class EmbeddingLoss:
- """Metric loss for training encoders to produce information-rich latent embeddings."""
-
- def __init__(self, margin: float = 0.2, type_of_triplets: str = "semihard") -> None:
- self.distance = distances.CosineSimilarity()
- self.reducer = reducers.ThresholdReducer(low=0)
- self.loss_fn = losses.TripletMarginLoss(
- margin=margin, distance=self.distance, reducer=self.reducer
- )
- self.miner = miners.MultiSimilarityMiner(epsilon=margin, distance=self.distance)
-
- def __call__(self, embeddings: Tensor, labels: Tensor) -> Tensor:
- """Computes the metric loss for the embeddings based on their labels.
-
- Args:
- embeddings (Tensor): The laten vectors encoded by the network.
- labels (Tensor): Labels of the embeddings.
-
- Returns:
- Tensor: The metric loss for the embeddings.
-
- """
- hard_pairs = self.miner(embeddings, labels)
- loss = self.loss_fn(embeddings, labels, hard_pairs)
- return loss
+__all__ = ["LabelSmoothingCrossEntropy"]
class LabelSmoothingCrossEntropy(nn.Module):
diff --git a/text_recognizer/networks/util.py b/text_recognizer/networks/util.py
index 131a6b4..d292680 100644
--- a/text_recognizer/networks/util.py
+++ b/text_recognizer/networks/util.py
@@ -1,38 +1,13 @@
"""Miscellaneous neural network functionality."""
import importlib
from pathlib import Path
-from typing import Dict, Tuple, Type
+from typing import Dict, Type
-from einops import rearrange
from loguru import logger
import torch
from torch import nn
-def sliding_window(
- images: torch.Tensor, patch_size: Tuple[int, int], stride: Tuple[int, int]
-) -> torch.Tensor:
- """Creates patches of an image.
-
- Args:
- images (torch.Tensor): A Torch tensor of a 4D image(s), i.e. (batch, channel, height, width).
- patch_size (Tuple[int, int]): The size of the patches to generate, e.g. 28x28 for EMNIST.
- stride (Tuple[int, int]): The stride of the sliding window.
-
- Returns:
- torch.Tensor: A tensor with the shape (batch, patches, height, width).
-
- """
- unfold = nn.Unfold(kernel_size=patch_size, stride=stride)
- # Preform the sliding window, unsqueeze as the channel dimesion is lost.
- c = images.shape[1]
- patches = unfold(images)
- patches = rearrange(
- patches, "b (c h w) t -> b t c h w", c=c, h=patch_size[0], w=patch_size[1],
- )
- return patches
-
-
def activation_function(activation: str) -> Type[nn.Module]:
"""Returns the callable activation function."""
activation_fns = nn.ModuleDict(
diff --git a/text_recognizer/networks/vq_transformer.py b/text_recognizer/networks/vq_transformer.py
deleted file mode 100644
index c673d96..0000000
--- a/text_recognizer/networks/vq_transformer.py
+++ /dev/null
@@ -1,150 +0,0 @@
-"""A VQ-Transformer for image to text recognition."""
-from typing import Dict, Optional, Tuple
-
-from einops import rearrange, repeat
-import torch
-from torch import nn
-from torch import Tensor
-
-from text_recognizer.networks.transformer import PositionalEncoding, Transformer
-from text_recognizer.networks.util import activation_function
-from text_recognizer.networks.util import configure_backbone
-from text_recognizer.networks.vqvae.encoder import _ResidualBlock
-
-
-class VQTransformer(nn.Module):
- """VQ+Transfomer for image to character sequence prediction."""
-
- def __init__(
- self,
- num_encoder_layers: int,
- num_decoder_layers: int,
- hidden_dim: int,
- vocab_size: int,
- num_heads: int,
- adaptive_pool_dim: Tuple,
- expansion_dim: int,
- dropout_rate: float,
- trg_pad_index: int,
- max_len: int,
- backbone: str,
- backbone_args: Optional[Dict] = None,
- activation: str = "gelu",
- ) -> None:
- super().__init__()
-
- # Configure vector quantized backbone.
- self.backbone = configure_backbone(backbone, backbone_args)
- self.conv = nn.Sequential(
- nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, stride=2),
- nn.ReLU(inplace=True),
- )
-
- # Configure embeddings for Transformer network.
- self.trg_pad_index = trg_pad_index
- self.vocab_size = vocab_size
- self.character_embedding = nn.Embedding(self.vocab_size, hidden_dim)
- self.src_position_embedding = nn.Parameter(torch.randn(1, max_len, hidden_dim))
- self.trg_position_encoding = PositionalEncoding(hidden_dim, dropout_rate)
- nn.init.normal_(self.character_embedding.weight, std=0.02)
-
- self.adaptive_pool = (
- nn.AdaptiveAvgPool2d((adaptive_pool_dim)) if adaptive_pool_dim else None
- )
-
- self.transformer = Transformer(
- num_encoder_layers,
- num_decoder_layers,
- hidden_dim,
- num_heads,
- expansion_dim,
- dropout_rate,
- activation,
- )
-
- self.head = nn.Sequential(nn.Linear(hidden_dim, vocab_size),)
-
- def _create_trg_mask(self, trg: Tensor) -> Tensor:
- # Move this outside the transformer.
- trg_pad_mask = (trg != self.trg_pad_index)[:, None, None]
- trg_len = trg.shape[1]
- trg_sub_mask = torch.tril(
- torch.ones((trg_len, trg_len), device=trg.device)
- ).bool()
- trg_mask = trg_pad_mask & trg_sub_mask
- return trg_mask
-
- def encoder(self, src: Tensor) -> Tensor:
- """Forward pass with the encoder of the transformer."""
- return self.transformer.encoder(src)
-
- def decoder(self, trg: Tensor, memory: Tensor, trg_mask: Tensor) -> Tensor:
- """Forward pass with the decoder of the transformer + classification head."""
- return self.head(
- self.transformer.decoder(trg=trg, memory=memory, trg_mask=trg_mask)
- )
-
- def extract_image_features(self, src: Tensor) -> Tuple[Tensor, Tensor]:
- """Extracts image features with a backbone neural network.
-
- It seem like the winning idea was to swap channels and width dimension and collapse
- the height dimension. The transformer is learning like a baby with this implementation!!! :D
- Ohhhh, the joy I am experiencing right now!! Bring in the beers! :D :D :D
-
- Args:
- src (Tensor): Input tensor.
-
- Returns:
- Tensor: The input src to the transformer and the vq loss.
-
- """
- # If batch dimension is missing, it needs to be added.
- if len(src.shape) < 4:
- src = src[(None,) * (4 - len(src.shape))]
- src, vq_loss = self.backbone.encode(src)
- # src = self.backbone.decoder.res_block(src)
- src = self.conv(src)
-
- if self.adaptive_pool is not None:
- src = rearrange(src, "b c h w -> b w c h")
- src = self.adaptive_pool(src)
- src = src.squeeze(3)
- else:
- src = rearrange(src, "b c h w -> b (w h) c")
-
- b, t, _ = src.shape
-
- src += self.src_position_embedding[:, :t]
-
- return src, vq_loss
-
- def target_embedding(self, trg: Tensor) -> Tensor:
- """Encodes target tensor with embedding and postion.
-
- Args:
- trg (Tensor): Target tensor.
-
- Returns:
- Tensor: Encoded target tensor.
-
- """
- trg = self.character_embedding(trg.long())
- trg = self.trg_position_encoding(trg)
- return trg
-
- def decode_image_features(
- self, image_features: Tensor, trg: Optional[Tensor] = None
- ) -> Tensor:
- """Takes images features from the backbone and decodes them with the transformer."""
- trg_mask = self._create_trg_mask(trg)
- trg = self.target_embedding(trg)
- out = self.transformer(image_features, trg, trg_mask=trg_mask)
-
- logits = self.head(out)
- return logits
-
- def forward(self, x: Tensor, trg: Optional[Tensor] = None) -> Tensor:
- """Forward pass with CNN transfomer."""
- image_features, vq_loss = self.extract_image_features(x)
- logits = self.decode_image_features(image_features, trg)
- return logits, vq_loss
generated by cgit v1.2.3-70-g09d2 (git 2.46.0) at 2024-12-25 18:52:49 +0000