Add new transformer network

4 files changed, 174 insertions, 4 deletions

diff --git a/text_recognizer/networks/image_transformer.py b/text_recognizer/networks/image_transformer.py
new file mode 100644
index 0000000..5a093dc
--- /dev/null
+++ b/text_recognizer/networks/image_transformer.py
@@ -0,0 +1,159 @@
+"""A Transformer with a cnn backbone.
+
+The network encodes a image with a convolutional backbone to a latent representation,
+i.e. feature maps. A 2d positional encoding is applied to the feature maps for 
+spatial information. The resulting feature are then set to a transformer decoder
+together with the target tokens. 
+
+TODO: Local attention for transformer.j
+
+"""
+import math
+from typing import Any, Dict, List, Optional, Sequence, Type
+
+from einops import rearrange
+import torch
+from torch import nn
+from torch import Tensor
+import torchvision
+
+from text_recognizer.data.emnist import emnist_mapping
+from text_recognizer.networks.transformer import (
+    Decoder,
+    DecoderLayer,
+    PositionalEncoding,
+    PositionalEncoding2D,
+    target_padding_mask,
+)
+
+
+class ImageTransformer(nn.Module):
+    def __init__(
+        self,
+        input_shape: Sequence[int],
+        output_shape: Sequence[int],
+        backbone: Type[nn.Module],
+        mapping: Optional[List[str]] = None,
+        num_decoder_layers: int = 4,
+        hidden_dim: int = 256,
+        num_heads: int = 4,
+        expansion_dim: int = 4,
+        dropout_rate: float = 0.1,
+        transformer_activation: str = "glu",
+    ) -> None:
+        # Configure mapping
+        mapping, inverse_mapping = self._configure_mapping(mapping)
+        self.vocab_size = len(mapping)
+        self.hidden_dim = hidden_dim
+        self.max_output_length = output_shape[0]
+        self.start_index = inverse_mapping["<s>"]
+        self.end_index = inverse_mapping["<e>"]
+        self.pad_index = inverse_mapping["<p>"]
+
+        # Image backbone
+        self.backbone = backbone
+        self.latent_encoding = PositionalEncoding2D(hidden_dim=hidden_dim, max_h=input_shape[1], max_w=input_shape[2])
+        
+        # Target token embedding
+        self.trg_embedding = nn.Embedding(self.vocab_size, hidden_dim)
+        self.trg_position_encoding = PositionalEncoding(hidden_dim, dropout_rate)
+
+        # Transformer decoder
+        self.decoder = Decoder(
+            decoder_layer=DecoderLayer(
+                hidden_dim=hidden_dim,
+                num_heads=num_heads,
+                expansion_dim=expansion_dim,
+                dropout_rate=dropout_rate,
+                activation=transformer_activation,
+            ),
+            num_layers=num_decoder_layers,
+            norm=nn.LayerNorm(hidden_dim),
+        )
+
+        # Classification head
+        self.head = nn.Linear(hidden_dim, self.vocab_size)
+
+        # Initialize weights
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        """Initialize network weights."""
+        self.trg_embedding.weight.data.uniform_(-0.1, 0.1)
+        self.head.bias.data.zero_()
+        self.head.weight.data.uniform_(-0.1, 0.1)
+
+        nn.init.kaiming_normal_(self.latent_encoding.weight.data, a=0, mode="fan_out", nonlinearity="relu")
+        if self.latent_encoding.bias is not None:
+            _, fan_out = nn.init._calculate_fan_in_and_fan_out(self.latent_encoding.weight.data)
+            bound = 1 / math.sqrt(fan_out)
+            nn.init.normal_(self.latent_encoding.bias, -bound, bound)
+
+    def _configure_mapping(self, mapping: Optional[List[str]]) -> Tuple[List[str], Dict[str, int]]:
+        """Configures mapping."""
+        if mapping is None:
+            mapping, inverse_mapping, _ = emnist_mapping() 
+        return mapping, inverse_mapping
+
+    def encode(self, image: Tensor) -> Tensor:
+        """Extracts image features with backbone.
+
+        Args:
+            image (Tensor): Image(s) of handwritten text.
+
+        Retuns:
+            Tensor: Image features.
+
+        Shapes:
+            - image: :math: `(B, C, H, W)` 
+            - latent: :math: `(B, T, C)`
+
+        """
+        # Extract image features.
+        latent = self.backbone(image)
+
+        # Add 2d encoding to the feature maps.
+        latent = self.latent_encoding(latent)
+        
+        # Collapse features maps height and width.
+        latent = rearrange(latent, "b c h w -> b (h w) c")
+        return latent
+
+    def decode(self, memory: Tensor, trg: Tensor) -> Tensor:
+        """Decodes image features with transformer decoder."""
+        trg_mask = target_padding_mask(trg=trg, pad_index=self.pad_index)
+        trg = self.trg_embedding(trg) * math.sqrt(self.hidden_dim)
+        trg = self.trg_position_encoding(trg)
+        out = self.decoder(trg=trg, memory=memory, trg_mask=trg_mask, memory_mask=None)
+        logits = self.head(out)
+        return logits
+
+    def predict(self, image: Tensor) -> Tensor:
+        """Transcribes text in image(s)."""
+        bsz = image.shape[0]
+        image_features = self.encode(image)
+
+        output_tokens = (torch.ones((bsz, self.max_output_length)) * self.pad_index).type_as(image).long()
+        output_tokens[:, 0] = self.start_index
+        for i in range(1, self.max_output_length):
+            trg = output_tokens[:, :i]
+            output = self.decode(image_features, trg)
+            output = torch.argmax(output, dim=-1)
+            output_tokens[:, i] = output[-1:]
+
+        # Set all tokens after end token to be padding.
+        for i in range(1, self.max_output_length):
+            indices = (output_tokens[:, i - 1] == self.end_index | (output_tokens[:, i - 1] == self.pad_index))
+            output_tokens[indices, i] = self.pad_index
+
+        return output_tokens
+
+
+
+
+
+
+
+
+
+
diff --git a/text_recognizer/networks/transformer/__init__.py b/text_recognizer/networks/transformer/__init__.py
index 9febc88..139cd23 100644
--- a/text_recognizer/networks/transformer/__init__.py
+++ b/text_recognizer/networks/transformer/__init__.py
@@ -1,3 +1,3 @@
 """Transformer modules."""
-from .positional_encoding import PositionalEncoding
+from .positional_encoding import PositionalEncoding, PositionalEncoding2D, target_padding_mask
 from .transformer import Decoder, Encoder, EncoderLayer, Transformer
diff --git a/text_recognizer/networks/transformer/attention.py b/text_recognizer/networks/transformer/attention.py
index cce1ecc..ac75d2f 100644
--- a/text_recognizer/networks/transformer/attention.py
+++ b/text_recognizer/networks/transformer/attention.py
@@ -50,8 +50,9 @@ class MultiHeadAttention(nn.Module):
         )
         nn.init.xavier_normal_(self.fc_out.weight)
 
+    @staticmethod
     def scaled_dot_product_attention(
-        self, query: Tensor, key: Tensor, value: Tensor, mask: Optional[Tensor] = None
+        query: Tensor, key: Tensor, value: Tensor, mask: Optional[Tensor] = None
     ) -> Tensor:
         """Calculates the scaled dot product attention."""
 
diff --git a/text_recognizer/networks/transformer/positional_encoding.py b/text_recognizer/networks/transformer/positional_encoding.py
index d67d297..dbde887 100644
--- a/text_recognizer/networks/transformer/positional_encoding.py
+++ b/text_recognizer/networks/transformer/positional_encoding.py
@@ -56,9 +56,9 @@ class PositionalEncoding2D(nn.Module):
         pe_h = repeat(pe_h, "h w d -> d h (w tile)", tile=max_w)
 
         pe_w = PositionalEncoding.make_pe(
-            hidden_dim // 2, max_len=max_h
+            hidden_dim // 2, max_len=max_w
         )  # [W, 1, D // 2]
-        pe_w = repeat(pe_w, "h w d -> d (h tile) w", tile=max_h)
+        pe_w = repeat(pe_w, "w h d -> d (h tile) w", tile=max_h)
 
         pe = torch.cat([pe_h, pe_w], dim=0)  # [D, H, W]
         return pe
@@ -70,3 +70,13 @@ class PositionalEncoding2D(nn.Module):
             raise ValueError("Hidden dimensions does not match.")
         x += self.pe[:, : x.shape[2], : x.shape[3]]
         return x
+
+def target_padding_mask(trg: Tensor, pad_index: int) -> Tensor:
+    """Returns causal target mask."""
+    trg_pad_mask = (trg != pad_index)[:, None, None]
+    trg_len = trg.shape[1]
+    trg_sub_mask = torch.tril(
+        torch.ones((trg_len, trg_len), device=trg.device)
+    ).bool()
+    trg_mask = trg_pad_mask & trg_sub_mask
+    return trg_mask