blob: d9782e828817f36c655e7bc10b90a01dd80318f5 (
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
|
"""Conformer block."""
from copy import deepcopy
from typing import Optional
from torch import nn, Tensor
from text_recognizer.networks.conformer.conv import ConformerConv
from text_recognizer.networks.conformer.mlp import MLP
from text_recognizer.networks.conformer.scale import Scale
from text_recognizer.networks.transformer.attention import Attention
from text_recognizer.networks.transformer.norm import PreNorm
class ConformerBlock(nn.Module):
def __init__(
self,
dim: int,
ff: MLP,
attn: Attention,
conv: ConformerConv,
) -> None:
super().__init__()
self.attn = PreNorm(dim, attn)
self.ff_1 = Scale(0.5, ff)
self.ff_2 = deepcopy(self.ff_1)
self.conv = conv
self.post_norm = nn.LayerNorm(dim)
def forward(self, x: Tensor, mask: Optional[Tensor]) -> Tensor:
x = self.ff_1(x) + x
x = self.attn(x, mask=mask) + x
x = self.conv(x) + x
x = self.ff_2(x) + x
return self.post_norm(x)
|
