# 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