summaryrefslogtreecommitdiff
path: root/text_recognizer/networks/transformer/decoder.py
blob: 1812e4053d92d933e98cf8c121465b1b9f136179 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
"""Transformer decoder module."""
from copy import deepcopy
from typing import Optional, Type

from torch import nn, Tensor

from text_recognizer.networks.transformer.attention import Attention
from text_recognizer.networks.transformer.mlp import FeedForward


class DecoderBlock(nn.Module):
    """Decoder block."""

    def __init__(
        self,
        self_attn: Attention,
        norm: Type[nn.Module],
        ff: FeedForward,
        cross_attn: Optional[Attention] = None,
    ) -> None:
        super().__init__()
        self.layers = ("self_attn", "cross_attn", "ff")
        self.blocks = self._build(self_attn, norm, ff, cross_attn)

    def _build(
        self,
        self_attn: Attention,
        norm: Type[nn.Module],
        ff: FeedForward,
        cross_attn: Optional[Attention],
    ) -> nn.ModuleDict:
        return nn.ModuleDict(
            {
                self.layers[0]: nn.ModuleList([norm, self_attn]),
                self.layers[1]: nn.ModuleList([deepcopy(norm), cross_attn]),
                self.layers[2]: nn.ModuleList([deepcopy(norm), ff]),
            }
        )

    def _apply_block(
        self,
        layer: str,
        x: Tensor,
        context: Optional[Tensor] = None,
        input_mask: Optional[Tensor] = None,
        context_mask: Optional[Tensor] = None,
    ) -> Tensor:
        """Applies block function."""
        residual = x
        norm_fn, layer_fn = self.blocks[layer]
        if layer == "self_attn":
            out = layer_fn(x=x, input_mask=input_mask)
        elif layer == "cross_attn":
            out = layer_fn(
                x=x, context=context, input_mask=input_mask, context_mask=context_mask
            )
        else:
            out = layer_fn(x)
        out += residual
        return norm_fn(out)

    def forward(
        self,
        x: Tensor,
        context: Optional[Tensor] = None,
        input_mask: Optional[Tensor] = None,
        context_mask: Optional[Tensor] = None,
    ) -> Tensor:
        """Applies decoder block on input signals."""
        for layer in self.layers:
            x = self._apply_block(
                layer=layer,
                x=x,
                context=context,
                input_mask=input_mask,
                context_mask=context_mask,
            )
        return x


class Decoder(nn.Module):
    """Decoder Network."""

    def __init__(self, depth: int, has_pos_emb: bool, block: DecoderBlock) -> None:
        super().__init__()
        self.depth = depth
        self.has_pos_emb = has_pos_emb
        self.blocks = nn.ModuleList([deepcopy(block) for _ in range(self.depth)])

    def forward(
        self,
        x: Tensor,
        context: Optional[Tensor] = None,
        input_mask: Optional[Tensor] = None,
        context_mask: Optional[Tensor] = None,
    ) -> Tensor:
        """Applies the network to the signals."""
        for block in self.blocks:
            x = block(
                x=x, context=context, input_mask=input_mask, context_mask=context_mask
            )
        return x