blob: c7da2265f216548024c7e0b6b9a30aaad020bee1 (
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
|
"""Transformer decoder module."""
from copy import deepcopy
from typing import Optional, Type
from torch import Tensor, nn
from text_recognizer.networks.transformer.attention import Attention
from text_recognizer.networks.transformer.decoder_block import DecoderBlock
from text_recognizer.networks.transformer.ff import FeedForward
class Decoder(nn.Module):
"""Decoder Network."""
def __init__(self, depth: int, dim: int, block: DecoderBlock) -> None:
super().__init__()
self.depth = depth
self.has_pos_emb = block.has_pos_emb
self.blocks = nn.ModuleList([deepcopy(block) for _ in range(self.depth)])
self.ln = nn.LayerNorm(dim)
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 self.ln(x)
|
