Repository: mindee/doctr

Issue:
[Models] Add SAR teacher forçing during training (decoder)
Teacher forçing (feeding the LSTM decoder of SAR with ground-truth characters during training) is not implemented, and it should really improve performances


Discussion:
Could you elaborate on the part of the paper where this is mentioned please? :pray: 
 "All the LSTM inputs are represented by one-hot vectors, followed by a linear ransformationΨ(). During training, the inputs of decoder LSTMs are replaced by the ground-truth character sequence"

Resolving patch:
diff --git a/doctr/models/recognition/sar.py b/doctr/models/recognition/sar.py
--- a/doctr/models/recognition/sar.py
+++ b/doctr/models/recognition/sar.py
@@ -89,7 +89,6 @@ class SARDecoder(layers.Layer, NestedObject):
         attention_units: number of hidden attention units
         num_decoder_layers: number of LSTM layers to stack
 
-
     """
     def __init__(
         self,
@@ -98,7 +97,7 @@ def __init__(
         vocab_size: int,
         embedding_units: int,
         attention_units: int,
-        num_decoder_layers: int = 2
+        num_decoder_layers: int = 2,
     ) -> None:
 
         super().__init__()
@@ -115,6 +114,7 @@ def call(
         self,
         features: tf.Tensor,
         holistic: tf.Tensor,
+        labels: Optional[tf.sparse.SparseTensor] = None,
         **kwargs: Any,
     ) -> tf.Tensor:
 
@@ -128,7 +128,7 @@ def call(
         # Initialize with the index of virtual START symbol (placed after <eos>)
         symbol = tf.fill(features.shape[0], self.vocab_size + 1)
         logits_list = []
-        for _ in range(self.max_length + 1):  # keep 1 step for <eos>
+        for t in range(self.max_length + 1):  # keep 1 step for <eos>
             # one-hot symbol with depth vocab_size + 1
             # embeded_symbol: shape (N, embedding_units)
             embeded_symbol = self.embed(tf.one_hot(symbol, depth=self.vocab_size + 1), **kwargs)
@@ -141,7 +141,13 @@ def call(
             # shape (N, rnn_units + 1) -> (N, vocab_size + 1)
             logits = self.output_dense(logits, **kwargs)
             # update symbol with predicted logits for t+1 step
-            symbol = tf.argmax(logits, axis=-1)
+            if kwargs.get('training'):
+                dense_labels = tf.sparse.to_dense(
+                    labels, default_value=self.vocab_size
+                )
+                symbol = dense_labels[:, t]
+            else:
+                symbol = tf.argmax(logits, axis=-1)
             logits_list.append(logits)
         outputs = tf.stack(logits_list, axis=1)  # shape (N, max_length + 1, vocab_size + 1)
 
@@ -196,13 +202,19 @@ def __init__(
     def call(
         self,
         x: tf.Tensor,
+        labels: Optional[tf.sparse.SparseTensor] = None,
         **kwargs: Any,
     ) -> tf.Tensor:
 
         features = self.feat_extractor(x, **kwargs)
         pooled_features = tf.reduce_max(features, axis=1)  # vertical max pooling
         encoded = self.encoder(pooled_features, **kwargs)
-        decoded = self.decoder(features, encoded, **kwargs)
+        if kwargs.get('training'):
+            if labels is None:
+                raise ValueError('Need to provide labels during training for teacher forcing')
+            decoded = self.decoder(features, encoded, labels, **kwargs)
+        else:
+            decoded = self.decoder(features, encoded, **kwargs)
 
         return decoded
 


Test patch:
diff --git a/test/test_models_recognition.py b/test/test_models_recognition.py
--- a/test/test_models_recognition.py
+++ b/test/test_models_recognition.py
@@ -52,6 +52,25 @@ def test_recognition_models(arch_name, input_shape, output_size):
     assert out.numpy().shape == (batch_size, *output_size)
 
 
+def test_sar_training():
+    batch_size = 4
+    input_shape = (64, 256, 3)
+    output_size = (41, 119)
+    reco_model = recognition.sar_vgg16_bn(input_shape=input_shape)
+    input_tensor = tf.random.uniform(shape=[batch_size, *input_shape], minval=0, maxval=1)
+    # input_labels: sparse_tensor of shape batch_size x max_len, encoding the labels
+    # filled with integers (classes of the characters at each timestep)
+    indices = [[0, 0], [0, 1], [1, 0], [1, 1], [1, 2], [2, 0], [3, 0], [3, 1], [3, 2], [3, 3], [3, 4]]
+    values = tf.random.uniform(shape=[11], minval=0, maxval=118, dtype=tf.dtypes.int64)
+    input_labels = tf.sparse.reorder(
+        tf.sparse.SparseTensor(indices=indices, values=values, dense_shape=[batch_size, 41])
+    )
+    out = reco_model(input_tensor, labels=input_labels, training=True)
+    assert isinstance(out, tf.Tensor)
+    assert isinstance(reco_model, tf.keras.Model)
+    assert out.numpy().shape == (batch_size, *output_size)
+
+
 @pytest.mark.parametrize(
     "post_processor, input_shape",
     [