Spaces:
Runtime error
Runtime error
| # Scene Text Recognition Model Hub | |
| # Copyright 2022 Darwin Bautista | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # https://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| import math | |
| from itertools import permutations | |
| from typing import Any, Optional, Sequence | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import Tensor | |
| from pytorch_lightning.utilities.types import STEP_OUTPUT | |
| from IndicPhotoOCR.utils.strhub.models.base import CrossEntropySystem | |
| from .model import PARSeq as Model | |
| class PARSeq(CrossEntropySystem): | |
| def __init__( | |
| self, | |
| charset_train: str, | |
| charset_test: str, | |
| max_label_length: int, | |
| batch_size: int, | |
| lr: float, | |
| warmup_pct: float, | |
| weight_decay: float, | |
| img_size: Sequence[int], | |
| patch_size: Sequence[int], | |
| embed_dim: int, | |
| enc_num_heads: int, | |
| enc_mlp_ratio: int, | |
| enc_depth: int, | |
| dec_num_heads: int, | |
| dec_mlp_ratio: int, | |
| dec_depth: int, | |
| perm_num: int, | |
| perm_forward: bool, | |
| perm_mirrored: bool, | |
| decode_ar: bool, | |
| refine_iters: int, | |
| dropout: float, | |
| **kwargs: Any, | |
| ) -> None: | |
| super().__init__(charset_train, charset_test, batch_size, lr, warmup_pct, weight_decay) | |
| self.save_hyperparameters() | |
| self.model = Model( | |
| len(self.tokenizer), | |
| max_label_length, | |
| img_size, | |
| patch_size, | |
| embed_dim, | |
| enc_num_heads, | |
| enc_mlp_ratio, | |
| enc_depth, | |
| dec_num_heads, | |
| dec_mlp_ratio, | |
| dec_depth, | |
| decode_ar, | |
| refine_iters, | |
| dropout, | |
| ) | |
| # Perm/attn mask stuff | |
| self.rng = np.random.default_rng() | |
| self.max_gen_perms = perm_num // 2 if perm_mirrored else perm_num | |
| self.perm_forward = perm_forward | |
| self.perm_mirrored = perm_mirrored | |
| def forward(self, images: Tensor, max_length: Optional[int] = None) -> Tensor: | |
| return self.model.forward(self.tokenizer, images, max_length) | |
| def gen_tgt_perms(self, tgt): | |
| """Generate shared permutations for the whole batch. | |
| This works because the same attention mask can be used for the shorter sequences | |
| because of the padding mask. | |
| """ | |
| # We don't permute the position of BOS, we permute EOS separately | |
| max_num_chars = tgt.shape[1] - 2 | |
| # Special handling for 1-character sequences | |
| if max_num_chars == 1: | |
| return torch.arange(3, device=self._device).unsqueeze(0) | |
| perms = [torch.arange(max_num_chars, device=self._device)] if self.perm_forward else [] | |
| # Additional permutations if needed | |
| max_perms = math.factorial(max_num_chars) | |
| if self.perm_mirrored: | |
| max_perms //= 2 | |
| num_gen_perms = min(self.max_gen_perms, max_perms) | |
| # For 4-char sequences and shorter, we generate all permutations and sample from the pool to avoid collisions | |
| # Note that this code path might NEVER get executed since the labels in a mini-batch typically exceed 4 chars. | |
| if max_num_chars < 5: | |
| # Pool of permutations to sample from. We only need the first half (if complementary option is selected) | |
| # Special handling for max_num_chars == 4 which correctly divides the pool into the flipped halves | |
| if max_num_chars == 4 and self.perm_mirrored: | |
| selector = [0, 3, 4, 6, 9, 10, 12, 16, 17, 18, 19, 21] | |
| else: | |
| selector = list(range(max_perms)) | |
| perm_pool = torch.as_tensor( | |
| list(permutations(range(max_num_chars), max_num_chars)), | |
| device=self._device, | |
| )[selector] | |
| # If the forward permutation is always selected, no need to add it to the pool for sampling | |
| if self.perm_forward: | |
| perm_pool = perm_pool[1:] | |
| perms = torch.stack(perms) | |
| if len(perm_pool): | |
| i = self.rng.choice(len(perm_pool), size=num_gen_perms - len(perms), replace=False) | |
| perms = torch.cat([perms, perm_pool[i]]) | |
| else: | |
| perms.extend( | |
| [torch.randperm(max_num_chars, device=self._device) for _ in range(num_gen_perms - len(perms))] | |
| ) | |
| perms = torch.stack(perms) | |
| if self.perm_mirrored: | |
| # Add complementary pairs | |
| comp = perms.flip(-1) | |
| # Stack in such a way that the pairs are next to each other. | |
| perms = torch.stack([perms, comp]).transpose(0, 1).reshape(-1, max_num_chars) | |
| # NOTE: | |
| # The only meaningful way of permuting the EOS position is by moving it one character position at a time. | |
| # However, since the number of permutations = T! and number of EOS positions = T + 1, the number of possible EOS | |
| # positions will always be much less than the number of permutations (unless a low perm_num is set). | |
| # Thus, it would be simpler to just train EOS using the full and null contexts rather than trying to evenly | |
| # distribute it across the chosen number of permutations. | |
| # Add position indices of BOS and EOS | |
| bos_idx = perms.new_zeros((len(perms), 1)) | |
| eos_idx = perms.new_full((len(perms), 1), max_num_chars + 1) | |
| perms = torch.cat([bos_idx, perms + 1, eos_idx], dim=1) | |
| # Special handling for the reverse direction. This does two things: | |
| # 1. Reverse context for the characters | |
| # 2. Null context for [EOS] (required for learning to predict [EOS] in NAR mode) | |
| if len(perms) > 1: | |
| perms[1, 1:] = max_num_chars + 1 - torch.arange(max_num_chars + 1, device=self._device) | |
| return perms | |
| def generate_attn_masks(self, perm): | |
| """Generate attention masks given a sequence permutation (includes pos. for bos and eos tokens) | |
| :param perm: the permutation sequence. i = 0 is always the BOS | |
| :return: lookahead attention masks | |
| """ | |
| sz = perm.shape[0] | |
| mask = torch.zeros((sz, sz), dtype=torch.bool, device=self._device) | |
| for i in range(sz): | |
| query_idx = perm[i] | |
| masked_keys = perm[i + 1 :] | |
| mask[query_idx, masked_keys] = True | |
| content_mask = mask[:-1, :-1].clone() | |
| mask[torch.eye(sz, dtype=torch.bool, device=self._device)] = True # mask "self" | |
| query_mask = mask[1:, :-1] | |
| return content_mask, query_mask | |
| def training_step(self, batch, batch_idx) -> STEP_OUTPUT: | |
| images, labels = batch | |
| tgt = self.tokenizer.encode(labels, self._device) | |
| # Encode the source sequence (i.e. the image codes) | |
| memory = self.model.encode(images) | |
| # Prepare the target sequences (input and output) | |
| tgt_perms = self.gen_tgt_perms(tgt) | |
| tgt_in = tgt[:, :-1] | |
| tgt_out = tgt[:, 1:] | |
| # The [EOS] token is not depended upon by any other token in any permutation ordering | |
| tgt_padding_mask = (tgt_in == self.pad_id) | (tgt_in == self.eos_id) | |
| loss = 0 | |
| loss_numel = 0 | |
| n = (tgt_out != self.pad_id).sum().item() | |
| for i, perm in enumerate(tgt_perms): | |
| tgt_mask, query_mask = self.generate_attn_masks(perm) | |
| out = self.model.decode(tgt_in, memory, tgt_mask, tgt_padding_mask, tgt_query_mask=query_mask) | |
| logits = self.model.head(out).flatten(end_dim=1) | |
| loss += n * F.cross_entropy(logits, tgt_out.flatten(), ignore_index=self.pad_id) | |
| loss_numel += n | |
| # After the second iteration (i.e. done with canonical and reverse orderings), | |
| # remove the [EOS] tokens for the succeeding perms | |
| if i == 1: | |
| tgt_out = torch.where(tgt_out == self.eos_id, self.pad_id, tgt_out) | |
| n = (tgt_out != self.pad_id).sum().item() | |
| loss /= loss_numel | |
| self.log('loss', loss) | |
| return loss | |