Repository: mindee/doctr

Issue:
[models] Enables dynamic sequence length
Currently, a lot of memory might be wasted if we do inference on a batch of sequences with length smaller than the fixed seq_len. Allowing a soft value `min(self.seq_len, max(len(s) for s in batch))` would save a lot of memory and computation.

But we need to make sure that this does not impact performances in a negative way first!


Discussion:

Resolving patch:
diff --git a/doctr/datasets/utils.py b/doctr/datasets/utils.py
--- a/doctr/datasets/utils.py
+++ b/doctr/datasets/utils.py
@@ -6,11 +6,12 @@
 import string
 import unicodedata
 import numpy as np
+from functools import partial
 from typing import List, Optional, Any
 
 from .vocabs import VOCABS
 
-__all__ = ['translate', 'encode_sequence', 'decode_sequence', 'encode_sequences']
+__all__ = ['translate', 'encode_string', 'decode_sequence', 'encode_sequences']
 
 
 def translate(
@@ -47,7 +48,7 @@ def translate(
     return translated
 
 
-def encode_sequence(
+def encode_string(
     input_string: str,
     vocab: str,
 ) -> List[int]:
@@ -74,12 +75,13 @@ def decode_sequence(
         mapping: vocabulary (string), the encoding is given by the indexing of the character sequence
 
     Returns:
-        A string, decoded from input_array"""
+        A string, decoded from input_array
+    """
 
     if not input_array.dtype == np.int_ or input_array.max() >= len(mapping):
         raise AssertionError("Input must be an array of int, with max less than mapping size")
-    decoded = ''.join(mapping[idx] for idx in input_array)
-    return decoded
+
+    return ''.join(map(mapping.__getitem__, input_array))
 
 
 def encode_sequences(
@@ -89,6 +91,7 @@ def encode_sequences(
     eos: int = -1,
     sos: Optional[int] = None,
     pad: Optional[int] = None,
+    dynamic_seq_length: bool = False,
     **kwargs: Any,
 ) -> np.ndarray:
     """Encode character sequences using a given vocab as mapping
@@ -100,6 +103,7 @@ def encode_sequences(
         eos: encoding of End Of String
         sos: optional encoding of Start Of String
         pad: optional encoding for padding. In case of padding, all sequences are followed by 1 EOS then PAD
+        dynamic_seq_length: if `target_size` is specified, uses it as upper bound and enables dynamic sequence size
 
     Returns:
         the padded encoded data as a tensor
@@ -108,29 +112,32 @@ def encode_sequences(
     if 0 <= eos < len(vocab):
         raise ValueError("argument 'eos' needs to be outside of vocab possible indices")
 
-    if not isinstance(target_size, int):
-        target_size = max(len(w) for w in sequences)
-        if sos:
-            target_size += 1
-        if pad:
-            target_size += 1
+    if not isinstance(target_size, int) or dynamic_seq_length:
+        # Maximum string length + EOS
+        max_length = max(len(w) for w in sequences) + 1
+        if isinstance(sos, int):
+            max_length += 1
+        if isinstance(pad, int):
+            max_length += 1
+        target_size = max_length if not isinstance(target_size, int) else min(max_length, target_size)
 
     # Pad all sequences
-    if pad:  # pad with padding symbol
+    if isinstance(pad, int):  # pad with padding symbol
         if 0 <= pad < len(vocab):
             raise ValueError("argument 'pad' needs to be outside of vocab possible indices")
         # In that case, add EOS at the end of the word before padding
-        encoded_data = np.full([len(sequences), target_size], pad, dtype=np.int32)
+        default_symbol = pad
     else:  # pad with eos symbol
-        encoded_data = np.full([len(sequences), target_size], eos, dtype=np.int32)
+        default_symbol = eos
+    encoded_data = np.full([len(sequences), target_size], default_symbol, dtype=np.int32)
 
-    for idx, seq in enumerate(sequences):
-        encoded_seq = encode_sequence(seq, vocab)
-        if pad:  # add eos at the end of the sequence
-            encoded_seq.append(eos)
-        encoded_data[idx, :min(len(encoded_seq), target_size)] = encoded_seq[:min(len(encoded_seq), target_size)]
+    # Encode the strings
+    for idx, seq in enumerate(map(partial(encode_string, vocab=vocab), sequences)):
+        if isinstance(pad, int):  # add eos at the end of the sequence
+            seq.append(eos)
+        encoded_data[idx, :min(len(seq), target_size)] = seq[:min(len(seq), target_size)]
 
-    if sos:  # place eos symbol at the beginning of each sequence
+    if isinstance(sos, int):  # place sos symbol at the beginning of each sequence
         if 0 <= sos < len(vocab):
             raise ValueError("argument 'sos' needs to be outside of vocab possible indices")
         encoded_data = np.roll(encoded_data, 1)


Test patch:
diff --git a/test/common/test_datasets_utils.py b/test/common/test_datasets_utils.py
--- a/test/common/test_datasets_utils.py
+++ b/test/common/test_datasets_utils.py
@@ -32,29 +32,30 @@ def test_translate(input_str, vocab, output_str):
 def test_encode_decode(input_str):
     mapping = """3K}7eé;5àÎYho]QwV6qU~W"XnbBvcADfËmy.9ÔpÛ*{CôïE%M4#ÈR:g@T$x?0î£|
         za1ù8,OG€P-kçHëÀÂ2É/ûIJ\'j(LNÙFut[)èZs+&°Sd=Ï!<â_Ç>rêi`l"""
-    encoded = utils.encode_sequence(input_str, mapping)
+    encoded = utils.encode_string(input_str, mapping)
     decoded = utils.decode_sequence(np.array(encoded), mapping)
     assert decoded == input_str
 
 
 @pytest.mark.parametrize(
-    "sequences, vocab, target_size, eos, sos, pad, error, out_shape, gts",
+    "sequences, vocab, target_size, sos, eos, pad, dynamic_len, error, out_shape, gts",
     [
-        [['cba'], 'abcdef', None, 1, None, None, True, (1, 3), [[2, 1, 0]]],
-        [['cba', 'a'], 'abcdef', None, -1, None, None, False, (2, 3), [[2, 1, 0], [0, -1, -1]]],
-        [['cba', 'a'], 'abcdef', None, 6, None, None, False, (2, 3), [[2, 1, 0], [0, 6, 6]]],
-        [['cba', 'a'], 'abcdef', 2, -1, None, None, False, (2, 2), [[2, 1], [0, -1]]],
-        [['cba', 'a'], 'abcdef', 4, -1, None, None, False, (2, 4), [[2, 1, 0, -1], [0, -1, -1, -1]]],
-        [['cba', 'a'], 'abcdef', 5, -1, 7, None, False, (2, 5), [[7, 2, 1, 0, -1], [7, 0, -1, -1, -1]]],
-        [['cba', 'a'], 'abcdef', None, -1, 7, 9, False, (2, 5), [[7, 2, 1, 0, -1], [7, 0, -1, 9, 9]]],
+        [['cba'], 'abcdef', None, None, 1, None, False, True, (1, 3), [[2, 1, 0]]],  # eos in vocab
+        [['cba', 'a'], 'abcdef', None, None, -1, None, False, False, (2, 4), [[2, 1, 0, -1], [0, -1, -1, -1]]],
+        [['cba', 'a'], 'abcdef', None, None, 6, None, False, False, (2, 4), [[2, 1, 0, 6], [0, 6, 6, 6]]],
+        [['cba', 'a'], 'abcdef', 2, None, -1, None, False, False, (2, 2), [[2, 1], [0, -1]]],
+        [['cba', 'a'], 'abcdef', 4, None, -1, None, False, False, (2, 4), [[2, 1, 0, -1], [0, -1, -1, -1]]],
+        [['cba', 'a'], 'abcdef', 5, 7, -1, None, False, False, (2, 5), [[7, 2, 1, 0, -1], [7, 0, -1, -1, -1]]],
+        [['cba', 'a'], 'abcdef', 6, 7, -1, None, True, False, (2, 5), [[7, 2, 1, 0, -1], [7, 0, -1, -1, -1]]],
+        [['cba', 'a'], 'abcdef', None, 7, -1, 9, False, False, (2, 6), [[7, 2, 1, 0, -1, 9], [7, 0, -1, 9, 9, 9]]],
     ],
 )
-def test_encode_sequences(sequences, vocab, target_size, eos, sos, pad, error, out_shape, gts):
+def test_encode_sequences(sequences, vocab, target_size, sos, eos, pad, dynamic_len, error, out_shape, gts):
     if error:
         with pytest.raises(ValueError):
-            _ = utils.encode_sequences(sequences, vocab, target_size, eos, sos, pad)
+            utils.encode_sequences(sequences, vocab, target_size, eos, sos, pad, dynamic_len)
     else:
-        out = utils.encode_sequences(sequences, vocab, target_size, eos, sos, pad)
+        out = utils.encode_sequences(sequences, vocab, target_size, eos, sos, pad, dynamic_len)
         assert isinstance(out, np.ndarray)
         assert out.shape == out_shape
         assert np.all(out == np.asarray(gts)), print(out, gts)