Repository: mindee/doctr

Issue:
[models] Implement HRGAN or MASTER
[This paper](https://arxiv.org/pdf/1904.01375v5.pdf) suggests a new architecture: Holistic Representation Guided Attention Network for text recognition model, inspired from transformers, which oustands SAR both in accuracy & speed.

We should implement this model, but the impressive speed results should be handled carefully (X8 speed compared to SAR), since the experiments are conducted on a GPU, and this model is highly parallelable (no recurrency). Is this new model so fast on CPU ?


Discussion:
[This one](https://arxiv.org/pdf/1910.02562.pdf) seems to be even stronger

Resolving patch:
diff --git a/doctr/models/recognition/__init__.py b/doctr/models/recognition/__init__.py
--- a/doctr/models/recognition/__init__.py
+++ b/doctr/models/recognition/__init__.py
@@ -2,3 +2,5 @@
 from .crnn import *
 from .sar import *
 from .zoo import *
+from .transformer import *
+from .master import *
diff --git a/doctr/models/recognition/master.py b/doctr/models/recognition/master.py
new file mode 100644
--- /dev/null
+++ b/doctr/models/recognition/master.py
@@ -0,0 +1,277 @@
+# Copyright (C) 2021, Mindee.
+
+# This program is licensed under the Apache License version 2.
+# See LICENSE or go to <https://www.apache.org/licenses/LICENSE-2.0.txt> for full license details.
+
+import tensorflow as tf
+from tensorflow.keras import layers, Sequential
+from typing import Tuple
+
+from .core import RecognitionModel
+from ..backbones.resnet import ResnetStage
+from ..utils import conv_sequence
+from .transformer import Decoder, positional_encoding, create_look_ahead_mask, create_padding_mask
+
+__all__ = ['MASTER']
+
+
+class MAGC(layers.Layer):
+
+    """Implements the Multi-Aspect Global Context Attention, as described in
+    <https://arxiv.org/pdf/1910.02562.pdf>`_.
+
+    Args:
+        inplanes: input channels
+        headers: number of headers to split channels
+        att_scale: if True, re-scale attention to counteract the variance distibutions
+        **kwargs
+    """
+
+    def __init__(
+        self,
+        inplanes: int,
+        headers: int = 1,
+        att_scale: bool = False,
+        **kwargs
+    ) -> None:
+        super().__init__(**kwargs)
+
+        self.headers = headers  # h
+        self.inplanes = inplanes  # C
+        self.att_scale = att_scale
+
+        self.single_header_inplanes = int(inplanes / headers)  # C / h
+
+        self.conv_mask = tf.keras.layers.Conv2D(
+            filters=1,
+            kernel_size=1,
+            kernel_initializer=tf.initializers.he_normal()
+        )
+
+        self.transform = tf.keras.Sequential(
+            [
+                tf.keras.layers.Conv2D(
+                    filters=self.inplanes,
+                    kernel_size=1,
+                    kernel_initializer=tf.initializers.he_normal()
+                ),
+                tf.keras.layers.LayerNormalization([1, 2, 3]),
+                tf.keras.layers.ReLU(),
+                tf.keras.layers.Conv2D(
+                    filters=self.inplanes,
+                    kernel_size=1,
+                    kernel_initializer=tf.initializers.he_normal()
+                ),
+            ],
+            name='transform'
+        )
+
+    @tf.function
+    def context_modeling(self, inputs: tf.Tensor) -> tf.Tensor:
+        b, h, w, c = (tf.shape(inputs)[i] for i in range(4))
+
+        # B, H, W, C -->> B*h, H, W, C/h
+        x = tf.reshape(inputs, shape=(b, h, w, self.headers, self.single_header_inplanes))
+        x = tf.transpose(x, perm=(0, 3, 1, 2, 4))
+        x = tf.reshape(x, shape=(b * self.headers, h, w, self.single_header_inplanes))
+
+        # Compute shorcut
+        shortcut = x
+        # B*h, 1, H*W, C/h
+        shortcut = tf.reshape(shortcut, shape=(b * self.headers, 1, h * w, self.single_header_inplanes))
+        # B*h, 1, C/h, H*W
+        shortcut = tf.transpose(shortcut, perm=[0, 1, 3, 2])
+
+        # Compute context mask
+        # B*h, H, W, 1,
+        context_mask = self.conv_mask(x)
+        # B*h, 1, H*W, 1
+        context_mask = tf.reshape(context_mask, shape=(b * self.headers, 1, h * w, 1))
+        # scale variance
+        if self.att_scale and self.headers > 1:
+            context_mask = context_mask / tf.sqrt(self.single_header_inplanes)
+        # B*h, 1, H*W, 1
+        context_mask = tf.keras.activations.softmax(context_mask, axis=2)
+
+        # Compute context
+        # B*h, 1, C/h, 1
+        context = tf.matmul(shortcut, context_mask)
+        context = tf.reshape(context, shape=(b, 1, c, 1))
+        # B, 1, 1, C
+        context = tf.transpose(context, perm=(0, 1, 3, 2))
+        # Set shape to resolve shape when calling this module in the Sequential MAGCResnet
+        batch, chan = inputs.get_shape().as_list()[0], inputs.get_shape().as_list()[-1]
+        context.set_shape([batch, 1, 1, chan])
+        return context
+
+    def call(self, inputs: tf.Tensor, **kwargs) -> tf.Tensor:
+        # Context modeling: B, H, W, C  ->  B, 1, 1, C
+        context = self.context_modeling(inputs)
+        # Transform: B, 1, 1, C  ->  B, 1, 1, C
+        transformed = self.transform(context)
+        return inputs + transformed
+
+
+class MAGCResnet(Sequential):
+
+    """Implements the modified resnet with MAGC layers, as described in paper.
+
+    Args:
+        headers: number of header to split channels in MAGC layers
+        input_shape: shape of the model input (without batch dim)
+    """
+
+    def __init__(
+        self,
+        headers: int = 1,
+        input_shape: Tuple[int, int, int] = (48, 160, 3),
+    ) -> None:
+        _layers = [
+            # conv_1x
+            *conv_sequence(out_channels=64, activation='relu', bn=True, kernel_size=3, input_shape=input_shape),
+            *conv_sequence(out_channels=128, activation='relu', bn=True, kernel_size=3),
+            layers.MaxPooling2D((2, 2), (2, 2)),
+            # conv_2x
+            ResnetStage(num_blocks=1, output_channels=256),
+            MAGC(inplanes=256, headers=headers, att_scale=True),
+            *conv_sequence(out_channels=256, activation='relu', bn=True, kernel_size=3),
+            layers.MaxPooling2D((2, 2), (2, 2)),
+            # conv_3x
+            ResnetStage(num_blocks=2, output_channels=512),
+            MAGC(inplanes=512, headers=headers, att_scale=True),
+            *conv_sequence(out_channels=512, activation='relu', bn=True, kernel_size=3),
+            layers.MaxPooling2D((2, 1), (2, 1)),
+            # conv_4x
+            ResnetStage(num_blocks=5, output_channels=512),
+            MAGC(inplanes=512, headers=headers, att_scale=True),
+            *conv_sequence(out_channels=512, activation='relu', bn=True, kernel_size=3),
+            # conv_5x
+            ResnetStage(num_blocks=3, output_channels=512),
+            MAGC(inplanes=512, headers=headers, att_scale=True),
+            *conv_sequence(out_channels=512, activation='relu', bn=True, kernel_size=3),
+        ]
+        super().__init__(_layers)
+
+
+class MASTER(RecognitionModel):
+
+    """Implements MASTER as described in paper: <https://arxiv.org/pdf/1910.02562.pdf>`_.
+    Implementation based on the official TF implementation: <https://github.com/jiangxiluning/MASTER-TF>`_.
+
+    Args:
+        vocab: vocabulary, (without EOS, SOS, PAD)
+        d_model: d parameter for the transformer decoder
+        headers: headers for the MAGC module
+        dff: depth of the pointwise feed-forward layer
+        num_heads: number of heads for the mutli-head attention module
+        num_layers: number of decoder layers to stack
+        max_length: maximum length of character sequence handled by the model
+        input_size: size of the image inputs
+    """
+
+    def __init__(
+        self,
+        vocab: str,
+        d_model: int = 512,
+        headers: int = 1,
+        dff: int = 2048,
+        num_heads: int = 8,
+        num_layers: int = 3,
+        max_length: int = 50,
+        input_size: tuple = (48, 160, 3),
+    ) -> None:
+        super().__init__(vocab=vocab)
+
+        self.input_size = input_size
+        self.max_length = max_length
+        self.vocab_size = len(vocab)
+
+        self.feature_extractor = MAGCResnet(headers=headers, input_shape=input_size)
+        self.seq_embedding = layers.Embedding(self.vocab_size + 1, d_model)  # One additional class for EOS
+
+        self.decoder = Decoder(
+            num_layers=num_layers,
+            d_model=d_model,
+            num_heads=num_heads,
+            dff=dff,
+            vocab_size=self.vocab_size,
+            maximum_position_encoding=max_length,
+        )
+        self.feature_pe = positional_encoding(input_size[0] * input_size[1], d_model)
+        self.linear = layers.Dense(self.vocab_size + 1, kernel_initializer=tf.initializers.he_uniform())
+
+    @tf.function
+    def make_mask(self, target: tf.Tensor) -> tf.Tensor:
+        look_ahead_mask = create_look_ahead_mask(tf.shape(target)[1])
+        target_padding_mask = create_padding_mask(target, self.vocab_size)  # Pad with EOS
+        combined_mask = tf.maximum(target_padding_mask, look_ahead_mask)
+        return combined_mask
+
+    def call(self, inputs: tf.Tensor, labels: tf.Tensor, **kwargs) -> tf.Tensor:
+        """Call function for training
+
+        Args:
+            inputs: images
+            labels: tensor of labels
+
+        Return:
+            Computed logits
+        """
+        # Encode
+        feature = self.feature_extractor(inputs, **kwargs)
+        b, h, w, c = (tf.shape(feature)[i] for i in range(4))
+        feature = tf.reshape(feature, shape=(b, h * w, c))
+        encoded = feature + self.feature_pe[:, :h * w, :]
+
+        tgt_mask = self.make_mask(labels)
+
+        output = self.decoder(labels, encoded, tgt_mask, None, training=True)
+        logits = self.linear(output)
+
+        return logits
+
+    @tf.function
+    def decode(self, inputs: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
+        """Decode function for prediction
+
+        Args:
+            inputs: images to predict
+
+        Return:
+            A Tuple of tf.Tensor: predictions, logits
+        """
+
+        feature = self.feature_extractor(inputs, training=False)
+        b, h, w, c = (tf.shape(feature)[i] for i in range(4))
+
+        feature = tf.reshape(feature, shape=(b, h * w, c))
+        encoded = feature + self.feature_pe[:, :h * w, :]
+
+        max_len = tf.constant(self.max_length, dtype=tf.int32)
+        start_symbol = tf.constant(self.vocab_size + 1, dtype=tf.int32)  # SOS (EOS = vocab_size)
+        padding_symbol = tf.constant(self.vocab_size, dtype=tf.int32)
+
+        ys = tf.fill(dims=(b, max_len - 1), value=padding_symbol)
+        start_vector = tf.fill(dims=(b, 1), value=start_symbol)
+        ys = tf.concat([start_vector, ys], axis=-1)
+
+        final_logits = tf.zeros(shape=(b, max_len - 1, self.vocab_size + 1), dtype=tf.float32)  # don't fgt EOS
+        # max_len = len + 2
+        for i in range(self.max_length - 1):
+            ys_mask = self.make_mask(ys)
+            output = self.decoder(ys, encoded, ys_mask, None, training=False)
+            logits = self.linear(output)
+            prob = tf.nn.softmax(logits, axis=-1)
+            next_word = tf.argmax(prob, axis=-1, output_type=ys.dtype)
+
+            # ys.shape = B, T
+            i_mesh, j_mesh = tf.meshgrid(tf.range(b), tf.range(max_len), indexing='ij')
+            indices = tf.stack([i_mesh[:, i + 1], j_mesh[:, i + 1]], axis=1)
+
+            ys = tf.tensor_scatter_nd_update(ys, indices, next_word[:, i + 1])
+
+            if i == (self.max_length - 2):
+                final_logits = logits
+
+        # ys predictions of shape B x max_length, final_logits of shape B x max_length x vocab_size + 1
+        return ys, final_logits
diff --git a/doctr/models/recognition/transformer.py b/doctr/models/recognition/transformer.py
new file mode 100644
--- /dev/null
+++ b/doctr/models/recognition/transformer.py
@@ -0,0 +1,240 @@
+# Copyright (C) 2021, Mindee.
+
+# This program is licensed under the Apache License version 2.
+# See LICENSE or go to <https://www.apache.org/licenses/LICENSE-2.0.txt> for full license details.
+
+# This module 'transformer.py' is 100% inspired from this Tensorflow tutorial:
+# https://www.tensorflow.org/text/tutorials/transformer
+
+
+from typing import Tuple
+
+import tensorflow as tf
+import numpy as np
+
+
+__all__ = ['Decoder', 'positional_encoding', 'create_look_ahead_mask', 'create_padding_mask']
+
+
+def get_angles(pos: np.array, i: np.array, d_model: int = 512) -> np.array:
+    """This function compute the 2D array of angles for sinusoidal positional encoding.
+
+    Args:
+        pos: range of positions to encode
+        i: range of depth to encode positions
+        d_model: depth parameter of the model
+
+    Returns:
+        2D array of angles, len(pos) x len(i)
+    """
+    angle_rates = 1 / np.power(10000, (2 * (i // 2)) / np.float32(d_model))
+    return pos * angle_rates
+
+
+def positional_encoding(position: int, d_model: int = 512) -> tf.Tensor:
+    """This function computes the 2D positional encoding of the position, on a depth d_model
+
+    Args:
+        position: Number of positions to encode
+        d_model: depth of the encoding
+
+    Returns:
+        2D positional encoding as described in Transformer paper.
+    """
+    angle_rads = get_angles(
+        np.arange(position)[:, np.newaxis],
+        np.arange(d_model)[np.newaxis, :],
+        d_model,
+    )
+    # apply sin to even indices in the array; 2i
+    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])
+    # apply cos to odd indices in the array; 2i+1
+    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])
+    pos_encoding = angle_rads[np.newaxis, ...]
+    return tf.cast(pos_encoding, dtype=tf.float32)
+
+
+@tf.function
+def create_padding_mask(seq: tf.Tensor, padding: int = 0) -> tf.Tensor:
+    seq = tf.cast(tf.math.equal(seq, padding), tf.float32)
+    # add extra dimensions to add the padding to the attention logits.
+    return seq[:, tf.newaxis, tf.newaxis, :]  # (batch_size, 1, 1, seq_len)
+
+
+@tf.function
+def create_look_ahead_mask(size: int) -> tf.Tensor:
+    mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)
+    return mask  # (seq_len, seq_len)
+
+
+@tf.function
+def scaled_dot_product_attention(
+    q: tf.Tensor, k: tf.Tensor, v: tf.Tensor, mask: tf.Tensor
+) -> Tuple[tf.Tensor, tf.Tensor]:
+
+    """Calculate the attention weights.
+    q, k, v must have matching leading dimensions.
+    k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.
+    The mask has different shapes depending on its type(padding or look ahead)
+    but it must be broadcastable for addition.
+    Args:
+        q: query shape == (..., seq_len_q, depth)
+        k: key shape == (..., seq_len_k, depth)
+        v: value shape == (..., seq_len_v, depth_v)
+        mask: Float tensor with shape broadcastable to (..., seq_len_q, seq_len_k). Defaults to None.
+    Returns:
+        output, attention_weights
+    """
+
+    matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)
+    # scale matmul_qk
+    dk = tf.cast(tf.shape(k)[-1], tf.float32)
+    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)
+    # add the mask to the scaled tensor.
+    if mask is not None:
+        scaled_attention_logits += (mask * -1e9)
+    # softmax is normalized on the last axis (seq_len_k) so that the scores
+    # add up to 1.
+    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)
+    output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)
+    return output
+
+
+class MultiHeadAttention(tf.keras.layers.Layer):
+
+    def __init__(self, d_model: int = 512, num_heads: int = 8) -> None:
+        super(MultiHeadAttention, self).__init__()
+        self.num_heads = num_heads
+        self.d_model = d_model
+
+        assert d_model % self.num_heads == 0
+
+        self.depth = d_model // self.num_heads
+
+        self.wq = tf.keras.layers.Dense(d_model, kernel_initializer=tf.initializers.he_uniform())
+        self.wk = tf.keras.layers.Dense(d_model, kernel_initializer=tf.initializers.he_uniform())
+        self.wv = tf.keras.layers.Dense(d_model, kernel_initializer=tf.initializers.he_uniform())
+
+        self.dense = tf.keras.layers.Dense(d_model, kernel_initializer=tf.initializers.he_uniform())
+
+    def split_heads(self, x: tf.Tensor, batch_size: int) -> tf.Tensor:
+        """Split the last dimension into (num_heads, depth).
+        Transpose the result such that the shape is (batch_size, num_heads, seq_len, depth)
+        """
+        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
+        return tf.transpose(x, perm=[0, 2, 1, 3])
+
+    def call(self, v: tf.Tensor, k: tf.Tensor, q: tf.Tensor, mask: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
+        batch_size = tf.shape(q)[0]
+
+        q = self.wq(q)  # (batch_size, seq_len, d_model)
+        k = self.wk(k)  # (batch_size, seq_len, d_model)
+        v = self.wv(v)  # (batch_size, seq_len, d_model)
+
+        q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)
+        k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)
+        v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)
+
+        # scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)
+        # attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)
+        scaled_attention = scaled_dot_product_attention(q, k, v, mask)
+
+        scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])  # (batch, seq_len_q, num_heads, depth)
+
+        concat_attention = tf.reshape(scaled_attention,
+                                      (batch_size, -1, self.d_model))  # (batch_size, seq_len_q, d_model)
+
+        output = self.dense(concat_attention)  # (batch_size, seq_len_q, d_model)
+
+        return output
+
+
+def point_wise_feed_forward_network(d_model: int = 512, dff: int = 2048) -> tf.keras.Sequential:
+    return tf.keras.Sequential([
+        tf.keras.layers.Dense(
+            dff, activation='relu', kernel_initializer=tf.initializers.he_uniform()
+        ),  # (batch, seq_len, dff)
+        tf.keras.layers.Dense(d_model, kernel_initializer=tf.initializers.he_uniform())  # (batch, seq_len, d_model)
+    ])
+
+
+class DecoderLayer(tf.keras.layers.Layer):
+
+    def __init__(self, d_model: int = 512, num_heads: int = 8, dff: int = 2048) -> None:
+        super(DecoderLayer, self).__init__()
+
+        self.mha1 = MultiHeadAttention(d_model, num_heads)
+        self.mha2 = MultiHeadAttention(d_model, num_heads)
+
+        self.ffn = point_wise_feed_forward_network(d_model, dff)
+
+        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
+        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
+        self.layernorm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
+
+    def call(
+        self,
+        x: tf.Tensor,
+        enc_output: tf.Tensor,
+        training: bool,
+        look_ahead_mask: tf.Tensor,
+        padding_mask: tf.Tensor
+    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
+        # enc_output.shape == (batch_size, input_seq_len, d_model)
+
+        attn1 = self.mha1(x, x, x, look_ahead_mask)  # (batch_size, target_seq_len, d_model)
+        out1 = self.layernorm1(attn1 + x)
+
+        attn2 = self.mha2(enc_output, enc_output, out1, padding_mask)  # (batch_size, target_seq_len, d_model)
+        out2 = self.layernorm2(attn2 + out1)  # (batch_size, target_seq_len, d_model)
+
+        ffn_output = self.ffn(out2)  # (batch_size, target_seq_len, d_model)
+        out3 = self.layernorm3(ffn_output + out2)  # (batch_size, target_seq_len, d_model)
+
+        return out3
+
+
+class Decoder(tf.keras.layers.Layer):
+
+    def __init__(
+        self,
+        num_layers: int = 3,
+        d_model: int = 512,
+        num_heads: int = 8,
+        dff: int = 2048,
+        vocab_size: int = 120,
+        maximum_position_encoding: int = 50,
+    ) -> None:
+        super(Decoder, self).__init__()
+
+        self.d_model = d_model
+        self.num_layers = num_layers
+
+        self.embedding = tf.keras.layers.Embedding(vocab_size + 2, d_model)  # 2 more classes EOS/SOS
+        self.pos_encoding = positional_encoding(maximum_position_encoding, d_model)
+
+        self.dec_layers = [DecoderLayer(d_model, num_heads, dff)
+                           for _ in range(num_layers)]
+
+    def call(
+        self,
+        x: tf.Tensor,
+        enc_output: tf.Tensor,
+        look_ahead_mask: tf.Tensor,
+        padding_mask: tf.Tensor,
+        training: bool = False,
+    ) -> Tuple[tf.Tensor, tf.Tensor]:
+
+        seq_len = tf.shape(x)[1]
+
+        x = self.embedding(x)  # (batch_size, target_seq_len, d_model)
+        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
+        x += self.pos_encoding[:, :seq_len, :]
+
+        for i in range(self.num_layers):
+            x = self.dec_layers[i](
+                x, enc_output, training, look_ahead_mask, padding_mask
+            )
+
+        # x.shape == (batch_size, target_seq_len, d_model)
+        return x


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
@@ -103,3 +103,19 @@ def test_recognition_zoo(arch_name):
 def test_recognition_zoo_error():
     with pytest.raises(ValueError):
         _ = recognition.zoo.recognition_predictor("my_fancy_model", pretrained=False)
+
+
+def test_master(mock_vocab, max_len=50, batch_size=16):
+    master = recognition.MASTER(vocab=mock_vocab, d_model=512, dff=512, num_heads=2, input_size=(48, 160, 3))
+    input_tensor = tf.random.uniform(shape=[batch_size, 48, 160, 3], minval=0, maxval=1)
+    mock_labels = tf.cast(
+        len(mock_vocab) * tf.random.uniform(shape=[batch_size, max_len], minval=0, maxval=1), tf.int32
+    )
+    logits = master(input_tensor, mock_labels)
+    assert isinstance(logits, tf.Tensor)
+    assert logits.shape == (batch_size, max_len, 1 + len(mock_vocab))  # 1 more for EOS
+    prediction, logits = master.decode(input_tensor)
+    assert isinstance(prediction, tf.Tensor)
+    assert isinstance(logits, tf.Tensor)
+    assert prediction.shape == (batch_size, max_len)
+    assert logits.shape == (batch_size, max_len, 1 + len(mock_vocab))