From 7268035fb9e57342612a8cc50a1fe04e8841ca2f Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Fri, 30 Jul 2021 23:15:03 +0200 Subject: attr bug fix, properly loading network --- text_recognizer/data/__init__.py | 6 -- text_recognizer/data/base_data_module.py | 2 +- text_recognizer/data/emnist_lines.py | 2 +- text_recognizer/data/iam_extended_paragraphs.py | 2 +- text_recognizer/data/iam_lines.py | 2 +- text_recognizer/data/iam_paragraphs.py | 2 +- text_recognizer/data/iam_preprocessor.py | 8 +- text_recognizer/data/iam_synthetic_paragraphs.py | 2 +- text_recognizer/data/mappings.py | 111 ++++++++++------------ text_recognizer/data/transforms.py | 16 ++-- text_recognizer/models/base.py | 4 +- text_recognizer/models/metrics.py | 8 +- text_recognizer/models/transformer.py | 80 ++++++++++++++-- text_recognizer/models/vqvae.py | 5 +- text_recognizer/networks/__init__.py | 4 - text_recognizer/networks/base.py | 18 ---- text_recognizer/networks/conv_transformer.py | 69 ++------------ text_recognizer/networks/transformer/attention.py | 2 +- text_recognizer/networks/transformer/layers.py | 16 ++-- text_recognizer/networks/transformer/norm.py | 8 +- text_recognizer/networks/util.py | 4 +- text_recognizer/networks/vqvae/vqvae.py | 1 - 22 files changed, 165 insertions(+), 207 deletions(-) delete mode 100644 text_recognizer/networks/base.py (limited to 'text_recognizer') diff --git a/text_recognizer/data/__init__.py b/text_recognizer/data/__init__.py index 3599a8b..2727b20 100644 --- a/text_recognizer/data/__init__.py +++ b/text_recognizer/data/__init__.py @@ -1,7 +1 @@ """Dataset modules.""" -from .base_dataset import BaseDataset, convert_strings_to_labels, split_dataset -from .base_data_module import BaseDataModule, load_and_print_info -from .download_utils import download_dataset -from .iam_paragraphs import IAMParagraphs -from .iam_synthetic_paragraphs import IAMSyntheticParagraphs -from .iam_extended_paragraphs import IAMExtendedParagraphs diff --git a/text_recognizer/data/base_data_module.py b/text_recognizer/data/base_data_module.py index 18b1996..408ae36 100644 --- a/text_recognizer/data/base_data_module.py +++ b/text_recognizer/data/base_data_module.py @@ -17,7 +17,7 @@ def load_and_print_info(data_module_class: type) -> None: print(dataset) -@attr.s +@attr.s(repr=False) class BaseDataModule(LightningDataModule): """Base PyTorch Lightning DataModule.""" diff --git a/text_recognizer/data/emnist_lines.py b/text_recognizer/data/emnist_lines.py index 4747508..7548ad5 100644 --- a/text_recognizer/data/emnist_lines.py +++ b/text_recognizer/data/emnist_lines.py @@ -32,7 +32,7 @@ IMAGE_X_PADDING = 28 MAX_OUTPUT_LENGTH = 89 # Same as IAMLines -@attr.s(auto_attribs=True) +@attr.s(auto_attribs=True, repr=False) class EMNISTLines(BaseDataModule): """EMNIST Lines dataset: synthetic handwritten lines dataset made from EMNIST,""" diff --git a/text_recognizer/data/iam_extended_paragraphs.py b/text_recognizer/data/iam_extended_paragraphs.py index 58c7369..23e424d 100644 --- a/text_recognizer/data/iam_extended_paragraphs.py +++ b/text_recognizer/data/iam_extended_paragraphs.py @@ -10,7 +10,7 @@ from text_recognizer.data.iam_paragraphs import IAMParagraphs from text_recognizer.data.iam_synthetic_paragraphs import IAMSyntheticParagraphs -@attr.s(auto_attribs=True) +@attr.s(auto_attribs=True, repr=False) class IAMExtendedParagraphs(BaseDataModule): augment: bool = attr.ib(default=True) diff --git a/text_recognizer/data/iam_lines.py b/text_recognizer/data/iam_lines.py index 13dd379..b7f3fdd 100644 --- a/text_recognizer/data/iam_lines.py +++ b/text_recognizer/data/iam_lines.py @@ -37,7 +37,7 @@ IMAGE_WIDTH = 1024 MAX_LABEL_LENGTH = 89 -@attr.s(auto_attribs=True) +@attr.s(auto_attribs=True, repr=False) class IAMLines(BaseDataModule): """IAM handwritten lines dataset.""" diff --git a/text_recognizer/data/iam_paragraphs.py b/text_recognizer/data/iam_paragraphs.py index de32875..82058e0 100644 --- a/text_recognizer/data/iam_paragraphs.py +++ b/text_recognizer/data/iam_paragraphs.py @@ -34,7 +34,7 @@ IMAGE_WIDTH = 1280 // IMAGE_SCALE_FACTOR MAX_LABEL_LENGTH = 682 -@attr.s(auto_attribs=True) +@attr.s(auto_attribs=True, repr=False) class IAMParagraphs(BaseDataModule): """IAM handwriting database paragraphs.""" diff --git a/text_recognizer/data/iam_preprocessor.py b/text_recognizer/data/iam_preprocessor.py index f7457e4..93a13bb 100644 --- a/text_recognizer/data/iam_preprocessor.py +++ b/text_recognizer/data/iam_preprocessor.py @@ -9,7 +9,7 @@ import collections import itertools from pathlib import Path import re -from typing import List, Optional, Union +from typing import List, Optional, Union, Sequence import click from loguru import logger @@ -57,15 +57,13 @@ class Preprocessor: lexicon_path: Optional[Union[str, Path]] = None, use_words: bool = False, prepend_wordsep: bool = False, - special_tokens: Optional[List[str]] = None, + special_tokens: Optional[Sequence[str]] = None, ) -> None: self.wordsep = "▁" self._use_word = use_words self._prepend_wordsep = prepend_wordsep self.special_tokens = special_tokens if special_tokens is not None else None - self.data_dir = Path(data_dir) - self.forms = load_metadata(self.data_dir, self.wordsep, use_words=use_words) # Load the set of graphemes: @@ -123,7 +121,7 @@ class Preprocessor: self.text.append(example["text"].lower()) def _to_index(self, line: str) -> torch.LongTensor: - if line in self.special_tokens: + if self.special_tokens is not None and line in self.special_tokens: return torch.LongTensor([self.tokens_to_index[line]]) token_to_index = self.graphemes_to_index if self.lexicon is not None: diff --git a/text_recognizer/data/iam_synthetic_paragraphs.py b/text_recognizer/data/iam_synthetic_paragraphs.py index a3697e7..f00a494 100644 --- a/text_recognizer/data/iam_synthetic_paragraphs.py +++ b/text_recognizer/data/iam_synthetic_paragraphs.py @@ -34,7 +34,7 @@ PROCESSED_DATA_DIRNAME = ( ) -@attr.s(auto_attribs=True) +@attr.s(auto_attribs=True, repr=False) class IAMSyntheticParagraphs(IAMParagraphs): """IAM Handwriting database of synthetic paragraphs.""" diff --git a/text_recognizer/data/mappings.py b/text_recognizer/data/mappings.py index 0d778b2..a934fd9 100644 --- a/text_recognizer/data/mappings.py +++ b/text_recognizer/data/mappings.py @@ -1,8 +1,9 @@ """Mapping to and from word pieces.""" from abc import ABC, abstractmethod from pathlib import Path -from typing import List, Optional, Union, Sequence +from typing import Dict, List, Optional, Union, Set, Sequence +import attr from loguru import logger import torch from torch import Tensor @@ -29,10 +30,17 @@ class AbstractMapping(ABC): ... +@attr.s class EmnistMapping(AbstractMapping): - def __init__(self, extra_symbols: Optional[Sequence[str]]) -> None: + extra_symbols: Optional[Set[str]] = attr.ib(default=None, converter=set) + mapping: Sequence[str] = attr.ib(init=False) + inverse_mapping: Dict[str, int] = attr.ib(init=False) + input_size: List[int] = attr.ib(init=False) + + def __attrs_post_init__(self) -> None: + """Post init configuration.""" self.mapping, self.inverse_mapping, self.input_size = emnist_mapping( - extra_symbols + self.extra_symbols ) def get_token(self, index: Union[int, Tensor]) -> str: @@ -54,42 +62,21 @@ class EmnistMapping(AbstractMapping): return Tensor([self.inverse_mapping[token] for token in text]) +@attr.s(auto_attribs=True) class WordPieceMapping(EmnistMapping): - def __init__( - self, - num_features: int = 1000, - tokens: str = "iamdb_1kwp_tokens_1000.txt", - lexicon: str = "iamdb_1kwp_lex_1000.txt", - data_dir: Optional[Union[str, Path]] = None, - use_words: bool = False, - prepend_wordsep: bool = False, - special_tokens: Sequence[str] = ("", "", "

"), - extra_symbols: Optional[Sequence[str]] = ("\n",), - ) -> None: - super().__init__(extra_symbols) - self.wordpiece_processor = self._configure_wordpiece_processor( - num_features, - tokens, - lexicon, - data_dir, - use_words, - prepend_wordsep, - special_tokens, - extra_symbols, - ) - - @staticmethod - def _configure_wordpiece_processor( - num_features: int, - tokens: str, - lexicon: str, - data_dir: Optional[Union[str, Path]], - use_words: bool, - prepend_wordsep: bool, - special_tokens: Optional[Sequence[str]], - extra_symbols: Optional[Sequence[str]], - ) -> Preprocessor: - data_dir = ( + data_dir: Optional[Path] = attr.ib(default=None) + num_features: int = attr.ib(default=1000) + tokens: str = attr.ib(default="iamdb_1kwp_tokens_1000.txt") + lexicon: str = attr.ib(default="iamdb_1kwp_lex_1000.txt") + use_words: bool = attr.ib(default=False) + prepend_wordsep: bool = attr.ib(default=False) + special_tokens: Set[str] = attr.ib(default={"", "", "

"}, converter=set) + extra_symbols: Set[str] = attr.ib(default={"\n",}, converter=set) + wordpiece_processor: Preprocessor = attr.ib(init=False) + + def __attrs_post_init__(self) -> None: + super().__attrs_post_init__() + self.data_dir = ( ( Path(__file__).resolve().parents[2] / "data" @@ -97,32 +84,32 @@ class WordPieceMapping(EmnistMapping): / "iam" / "iamdb" ) - if data_dir is None - else Path(data_dir) + if self.data_dir is None + else Path(self.data_dir) ) - - logger.debug(f"Using data dir: {data_dir}") - if not data_dir.exists(): - raise RuntimeError(f"Could not locate iamdb directory at {data_dir}") + logger.debug(f"Using data dir: {self.data_dir}") + if not self.data_dir.exists(): + raise RuntimeError(f"Could not locate iamdb directory at {self.data_dir}") processed_path = ( Path(__file__).resolve().parents[2] / "data" / "processed" / "iam_lines" ) - tokens_path = processed_path / tokens - lexicon_path = processed_path / lexicon - - if extra_symbols is not None: - special_tokens += extra_symbols - - return Preprocessor( - data_dir, - num_features, - tokens_path, - lexicon_path, - use_words, - prepend_wordsep, - special_tokens, + tokens_path = processed_path / self.tokens + lexicon_path = processed_path / self.lexicon + + special_tokens = self.special_tokens + if self.extra_symbols is not None: + special_tokens = special_tokens | self.extra_symbols + + self.wordpiece_processor = Preprocessor( + data_dir=self.data_dir, + num_features=self.num_features, + tokens_path=tokens_path, + lexicon_path=lexicon_path, + use_words=self.use_words, + prepend_wordsep=self.prepend_wordsep, + special_tokens=special_tokens, ) def __len__(self) -> int: @@ -151,7 +138,9 @@ class WordPieceMapping(EmnistMapping): text = text.lower().replace(" ", "▁") return torch.LongTensor(self.wordpiece_processor.to_index(text)) - def __getitem__(self, x: Union[str, int, Tensor]) -> Union[str, Tensor]: + def __getitem__(self, x: Union[str, int, List[int], Tensor]) -> Union[str, Tensor]: + if isinstance(x, int): + x = [x] if isinstance(x, str): - return self.get_index(x) - return self.get_token(x) + return self.get_indices(x) + return self.get_text(x) diff --git a/text_recognizer/data/transforms.py b/text_recognizer/data/transforms.py index 66531a5..3b1b929 100644 --- a/text_recognizer/data/transforms.py +++ b/text_recognizer/data/transforms.py @@ -24,14 +24,14 @@ class WordPiece: max_len: int = 451, ) -> None: self.mapping = WordPieceMapping( - num_features, - tokens, - lexicon, - data_dir, - use_words, - prepend_wordsep, - special_tokens, - extra_symbols, + data_dir=data_dir, + num_features=num_features, + tokens=tokens, + lexicon=lexicon, + use_words=use_words, + prepend_wordsep=prepend_wordsep, + special_tokens=special_tokens, + extra_symbols=extra_symbols, ) self.max_len = max_len diff --git a/text_recognizer/models/base.py b/text_recognizer/models/base.py index 3e02261..dfb4ca4 100644 --- a/text_recognizer/models/base.py +++ b/text_recognizer/models/base.py @@ -11,8 +11,6 @@ from torch import nn from torch import Tensor import torchmetrics -from text_recognizer.networks.base import BaseNetwork - @attr.s class BaseLitModel(LightningModule): @@ -21,7 +19,7 @@ class BaseLitModel(LightningModule): def __attrs_pre_init__(self) -> None: super().__init__() - network: Type[BaseNetwork] = attr.ib() + network: Type[nn.Module] = attr.ib() criterion_config: DictConfig = attr.ib(converter=DictConfig) optimizer_config: DictConfig = attr.ib(converter=DictConfig) lr_scheduler_config: DictConfig = attr.ib(converter=DictConfig) diff --git a/text_recognizer/models/metrics.py b/text_recognizer/models/metrics.py index 4117ae2..9793157 100644 --- a/text_recognizer/models/metrics.py +++ b/text_recognizer/models/metrics.py @@ -1,5 +1,5 @@ """Character Error Rate (CER).""" -from typing import Set, Sequence +from typing import Set import attr import editdistance @@ -12,7 +12,7 @@ from torchmetrics import Metric class CharacterErrorRate(Metric): """Character error rate metric, computed using Levenshtein distance.""" - ignore_tokens: Set = attr.ib(converter=set) + ignore_indices: Set = attr.ib(converter=set) error: Tensor = attr.ib(init=False) total: Tensor = attr.ib(init=False) @@ -25,8 +25,8 @@ class CharacterErrorRate(Metric): """Update CER.""" bsz = preds.shape[0] for index in range(bsz): - pred = [p for p in preds[index].tolist() if p not in self.ignore_tokens] - target = [t for t in targets[index].tolist() if t not in self.ignore_tokens] + pred = [p for p in preds[index].tolist() if p not in self.ignore_indices] + target = [t for t in targets[index].tolist() if t not in self.ignore_indices] distance = editdistance.distance(pred, target) error = distance / max(len(pred), len(target)) self.error += error diff --git a/text_recognizer/models/transformer.py b/text_recognizer/models/transformer.py index f5cb491..7a9d566 100644 --- a/text_recognizer/models/transformer.py +++ b/text_recognizer/models/transformer.py @@ -1,11 +1,11 @@ """PyTorch Lightning model for base Transformers.""" -from typing import Dict, List, Optional, Sequence, Union, Tuple, Type +from typing import Sequence, Tuple, Type import attr -import hydra -from omegaconf import DictConfig -from torch import nn, Tensor +import torch +from torch import Tensor +from text_recognizer.data.mappings import AbstractMapping from text_recognizer.models.metrics import CharacterErrorRate from text_recognizer.models.base import BaseLitModel @@ -13,18 +13,31 @@ from text_recognizer.models.base import BaseLitModel @attr.s(auto_attribs=True) class TransformerLitModel(BaseLitModel): """A PyTorch Lightning model for transformer networks.""" + mapping: Type[AbstractMapping] = attr.ib() + start_token: str = attr.ib() + end_token: str = attr.ib() + pad_token: str = attr.ib() - ignore_tokens: Sequence[str] = attr.ib(default=("", "", "

",)) + start_index: Tensor = attr.ib(init=False) + end_index: Tensor = attr.ib(init=False) + pad_index: Tensor = attr.ib(init=False) + + ignore_indices: Sequence[str] = attr.ib(init=False) val_cer: CharacterErrorRate = attr.ib(init=False) test_cer: CharacterErrorRate = attr.ib(init=False) def __attrs_post_init__(self) -> None: - self.val_cer = CharacterErrorRate(self.ignore_tokens) - self.test_cer = CharacterErrorRate(self.ignore_tokens) + """Post init configuration.""" + self.start_index = self.mapping.get_index(self.start_token) + self.end_index = self.mapping.get_index(self.end_token) + self.pad_index = self.mapping.get_index(self.pad_token) + self.ignore_indices = [self.start_index, self.end_index, self.pad_index] + self.val_cer = CharacterErrorRate(self.ignore_indices) + self.test_cer = CharacterErrorRate(self.ignore_indices) def forward(self, data: Tensor) -> Tensor: """Forward pass with the transformer network.""" - return self.network.predict(data) + return self.predict(data) def training_step(self, batch: Tuple[Tensor, Tensor], batch_idx: int) -> Tensor: """Training step.""" @@ -38,17 +51,64 @@ class TransformerLitModel(BaseLitModel): """Validation step.""" data, targets = batch + # Compute the loss. logits = self.network(data, targets[:-1]) loss = self.loss_fn(logits, targets[1:]) self.log("val/loss", loss, prog_bar=True) - pred = self.network.predict(data) + # Get the token prediction. + pred = self(data) self.val_cer(pred, targets) self.log("val/cer", self.val_cer, on_step=False, on_epoch=True, prog_bar=True) def test_step(self, batch: Tuple[Tensor, Tensor], batch_idx: int) -> None: """Test step.""" data, targets = batch - pred = self.network.predict(data) + + # Compute the text prediction. + pred = self(data) self.test_cer(pred, targets) self.log("test/cer", self.test_cer, on_step=False, on_epoch=True, prog_bar=True) + + def predict(self, x: Tensor) -> Tensor: + """Predicts text in image. + + Args: + x (Tensor): Image(s) to extract text from. + + Shapes: + - x: :math: `(B, H, W)` + - output: :math: `(B, S)` + + Returns: + Tensor: A tensor of token indices of the predictions from the model. + """ + bsz = x.shape[0] + + # Encode image(s) to latent vectors. + z = self.network.encode(x) + + # Create a placeholder matrix for storing outputs from the network + output = torch.ones((bsz, self.max_output_len), dtype=torch.long).to(x.device) + output[:, 0] = self.start_index + + for i in range(1, self.max_output_len): + context = output[:, :i] # (bsz, i) + logits = self.network.decode(z, context) # (i, bsz, c) + tokens = torch.argmax(logits, dim=-1) # (i, bsz) + output[:, i : i + 1] = tokens[-1:] + + # Early stopping of prediction loop if token is end or padding token. + if ( + output[:, i - 1] == self.end_index | output[: i - 1] == self.pad_index + ).all(): + break + + # Set all tokens after end token to pad token. + for i in range(1, self.max_output_len): + idx = ( + output[:, i - 1] == self.end_index | output[:, i - 1] == self.pad_index + ) + output[idx, i] = self.pad_index + + return output diff --git a/text_recognizer/models/vqvae.py b/text_recognizer/models/vqvae.py index 0172163..e215e14 100644 --- a/text_recognizer/models/vqvae.py +++ b/text_recognizer/models/vqvae.py @@ -34,8 +34,6 @@ class VQVAELitModel(BaseLitModel): loss = self.loss_fn(reconstructions, data) loss += vq_loss self.log("val/loss", loss, prog_bar=True) - title = "val_pred_examples" - self._log_prediction(data, reconstructions, title) def test_step(self, batch: Tuple[Tensor, Tensor], batch_idx: int) -> None: """Test step.""" @@ -43,5 +41,4 @@ class VQVAELitModel(BaseLitModel): reconstructions, vq_loss = self.network(data) loss = self.loss_fn(reconstructions, data) loss += vq_loss - title = "test_pred_examples" - self._log_prediction(data, reconstructions, title) + self.log("test/loss", loss) diff --git a/text_recognizer/networks/__init__.py b/text_recognizer/networks/__init__.py index 618450f..d9ef58b 100644 --- a/text_recognizer/networks/__init__.py +++ b/text_recognizer/networks/__init__.py @@ -1,5 +1 @@ """Network modules""" -# from .encoders import EfficientNet -from .vqvae import VQVAE - -# from .cnn_transformer import CNNTransformer diff --git a/text_recognizer/networks/base.py b/text_recognizer/networks/base.py deleted file mode 100644 index 07b6a32..0000000 --- a/text_recognizer/networks/base.py +++ /dev/null @@ -1,18 +0,0 @@ -"""Base network with required methods.""" -from abc import abstractmethod - -import attr -from torch import nn, Tensor - - -@attr.s -class BaseNetwork(nn.Module): - """Base network.""" - - def __attrs_pre_init__(self) -> None: - super().__init__() - - @abstractmethod - def predict(self, x: Tensor) -> Tensor: - """Return token indices for predictions.""" - ... diff --git a/text_recognizer/networks/conv_transformer.py b/text_recognizer/networks/conv_transformer.py index 4acdc36..7371be4 100644 --- a/text_recognizer/networks/conv_transformer.py +++ b/text_recognizer/networks/conv_transformer.py @@ -1,13 +1,10 @@ """Vision transformer for character recognition.""" import math -from typing import Tuple, Type +from typing import Tuple import attr -import torch from torch import nn, Tensor -from text_recognizer.data.mappings import AbstractMapping -from text_recognizer.networks.base import BaseNetwork from text_recognizer.networks.encoders.efficientnet import EfficientNet from text_recognizer.networks.transformer.layers import Decoder from text_recognizer.networks.transformer.positional_encodings import ( @@ -16,25 +13,24 @@ from text_recognizer.networks.transformer.positional_encodings import ( ) -@attr.s(auto_attribs=True) -class ConvTransformer(BaseNetwork): +@attr.s +class ConvTransformer(nn.Module): + """Convolutional encoder and transformer decoder network.""" + + def __attrs_pre_init__(self) -> None: + super().__init__() + # Parameters and placeholders, input_dims: Tuple[int, int, int] = attr.ib() hidden_dim: int = attr.ib() dropout_rate: float = attr.ib() max_output_len: int = attr.ib() num_classes: int = attr.ib() - start_token: str = attr.ib() - start_index: Tensor = attr.ib(init=False) - end_token: str = attr.ib() - end_index: Tensor = attr.ib(init=False) - pad_token: str = attr.ib() - pad_index: Tensor = attr.ib(init=False) + pad_index: Tensor = attr.ib() # Modules. encoder: EfficientNet = attr.ib() decoder: Decoder = attr.ib() - mapping: Type[AbstractMapping] = attr.ib() latent_encoder: nn.Sequential = attr.ib(init=False) token_embedding: nn.Embedding = attr.ib(init=False) @@ -43,10 +39,6 @@ class ConvTransformer(BaseNetwork): def __attrs_post_init__(self) -> None: """Post init configuration.""" - self.start_index = self.mapping.get_index(self.start_token) - self.end_index = self.mapping.get_index(self.end_token) - self.pad_index = self.mapping.get_index(self.pad_token) - # Latent projector for down sampling number of filters and 2d # positional encoding. self.latent_encoder = nn.Sequential( @@ -156,46 +148,3 @@ class ConvTransformer(BaseNetwork): z = self.encode(x) logits = self.decode(z, context) return logits - - def predict(self, x: Tensor) -> Tensor: - """Predicts text in image. - - Args: - x (Tensor): Image(s) to extract text from. - - Shapes: - - x: :math: `(B, H, W)` - - output: :math: `(B, S)` - - Returns: - Tensor: A tensor of token indices of the predictions from the model. - """ - bsz = x.shape[0] - - # Encode image(s) to latent vectors. - z = self.encode(x) - - # Create a placeholder matrix for storing outputs from the network - output = torch.ones((bsz, self.max_output_len), dtype=torch.long).to(x.device) - output[:, 0] = self.start_index - - for i in range(1, self.max_output_len): - context = output[:, :i] # (bsz, i) - logits = self.decode(z, context) # (i, bsz, c) - tokens = torch.argmax(logits, dim=-1) # (i, bsz) - output[:, i : i + 1] = tokens[-1:] - - # Early stopping of prediction loop if token is end or padding token. - if ( - output[:, i - 1] == self.end_index | output[: i - 1] == self.pad_index - ).all(): - break - - # Set all tokens after end token to pad token. - for i in range(1, self.max_output_len): - idx = ( - output[:, i - 1] == self.end_index | output[:, i - 1] == self.pad_index - ) - output[idx, i] = self.pad_index - - return output diff --git a/text_recognizer/networks/transformer/attention.py b/text_recognizer/networks/transformer/attention.py index 2770dc1..9202cce 100644 --- a/text_recognizer/networks/transformer/attention.py +++ b/text_recognizer/networks/transformer/attention.py @@ -24,9 +24,9 @@ class Attention(nn.Module): dim: int = attr.ib() num_heads: int = attr.ib() + causal: bool = attr.ib(default=False) dim_head: int = attr.ib(default=64) dropout_rate: float = attr.ib(default=0.0) - casual: bool = attr.ib(default=False) scale: float = attr.ib(init=False) dropout: nn.Dropout = attr.ib(init=False) fc: nn.Linear = attr.ib(init=False) diff --git a/text_recognizer/networks/transformer/layers.py b/text_recognizer/networks/transformer/layers.py index 9b2f236..66c9c50 100644 --- a/text_recognizer/networks/transformer/layers.py +++ b/text_recognizer/networks/transformer/layers.py @@ -30,8 +30,7 @@ class AttentionLayers(nn.Module): causal: bool = attr.ib(default=False) cross_attend: bool = attr.ib(default=False) pre_norm: bool = attr.ib(default=True) - rotary_emb: Optional[RotaryEmbedding] = attr.ib(default=None, init=False) - has_pos_emb: bool = attr.ib(init=False) + rotary_emb: Optional[RotaryEmbedding] = attr.ib(default=None) layer_types: Tuple[str, ...] = attr.ib(init=False) layers: nn.ModuleList = attr.ib(init=False) attn: partial = attr.ib(init=False) @@ -40,12 +39,11 @@ class AttentionLayers(nn.Module): def __attrs_post_init__(self) -> None: """Post init configuration.""" - self.has_pos_emb = True if self.rotary_emb is not None else False self.layer_types = self._get_layer_types() * self.depth attn = load_partial_fn( self.attn_fn, dim=self.dim, num_heads=self.num_heads, **self.attn_kwargs ) - norm = load_partial_fn(self.norm_fn, dim=self.dim) + norm = load_partial_fn(self.norm_fn, normalized_shape=self.dim) ff = load_partial_fn(self.ff_fn, dim=self.dim, **self.ff_kwargs) self.layers = self._build_network(attn, norm, ff) @@ -103,13 +101,11 @@ class AttentionLayers(nn.Module): return x +@attr.s(auto_attribs=True) class Encoder(AttentionLayers): - def __init__(self, **kwargs: Any) -> None: - assert "causal" not in kwargs, "Cannot set causality on encoder" - super().__init__(causal=False, **kwargs) + causal: bool = attr.ib(default=False, init=False) +@attr.s(auto_attribs=True) class Decoder(AttentionLayers): - def __init__(self, **kwargs: Any) -> None: - assert "causal" not in kwargs, "Cannot set causality on decoder" - super().__init__(causal=True, **kwargs) + causal: bool = attr.ib(default=True, init=False) diff --git a/text_recognizer/networks/transformer/norm.py b/text_recognizer/networks/transformer/norm.py index 8bc3221..4930adf 100644 --- a/text_recognizer/networks/transformer/norm.py +++ b/text_recognizer/networks/transformer/norm.py @@ -12,9 +12,9 @@ from torch import Tensor class ScaleNorm(nn.Module): - def __init__(self, dim: int, eps: float = 1.0e-5) -> None: + def __init__(self, normalized_shape: int, eps: float = 1.0e-5) -> None: super().__init__() - self.scale = dim ** -0.5 + self.scale = normalized_shape ** -0.5 self.eps = eps self.g = nn.Parameter(torch.ones(1)) @@ -24,9 +24,9 @@ class ScaleNorm(nn.Module): class PreNorm(nn.Module): - def __init__(self, dim: int, fn: Type[nn.Module]) -> None: + def __init__(self, normalized_shape: int, fn: Type[nn.Module]) -> None: super().__init__() - self.norm = nn.LayerNorm(dim) + self.norm = nn.LayerNorm(normalized_shape) self.fn = fn def forward(self, x: Tensor, **kwargs: Dict) -> Tensor: diff --git a/text_recognizer/networks/util.py b/text_recognizer/networks/util.py index e822c57..c94e8dc 100644 --- a/text_recognizer/networks/util.py +++ b/text_recognizer/networks/util.py @@ -24,6 +24,6 @@ def activation_function(activation: str) -> Type[nn.Module]: def load_partial_fn(fn: str, **kwargs: Any) -> partial: - """Loads partial function.""" + """Loads partial function/class.""" module = import_module(".".join(fn.split(".")[:-1])) - return partial(getattr(module, fn.split(".")[0]), **kwargs) + return partial(getattr(module, fn.split(".")[-1]), **kwargs) diff --git a/text_recognizer/networks/vqvae/vqvae.py b/text_recognizer/networks/vqvae/vqvae.py index 1f08e5e..5aa929b 100644 --- a/text_recognizer/networks/vqvae/vqvae.py +++ b/text_recognizer/networks/vqvae/vqvae.py @@ -1,5 +1,4 @@ """The VQ-VAE.""" - from typing import Any, Dict, List, Optional, Tuple from torch import nn -- cgit v1.2.3-70-g09d2