From 49ca6ade1a19f7f9c702171537fe4be0dfcda66d Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Fri, 25 Aug 2023 23:19:14 +0200 Subject: Rename and add flash atten --- text_recognizer/network/transformer/embedding/token.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) create mode 100644 text_recognizer/network/transformer/embedding/token.py (limited to 'text_recognizer/network/transformer/embedding/token.py') diff --git a/text_recognizer/network/transformer/embedding/token.py b/text_recognizer/network/transformer/embedding/token.py new file mode 100644 index 0000000..1df2fd6 --- /dev/null +++ b/text_recognizer/network/transformer/embedding/token.py @@ -0,0 +1,18 @@ +from torch import nn, Tensor + +from text_recognizer.network.transformer.embedding.l2_norm import l2_norm + + +class TokenEmbedding(nn.Module): + def __init__(self, num_tokens: int, dim: int, use_l2: bool = True) -> None: + super().__init__() + self.use_l2 = use_l2 + self.to_embedding = nn.Embedding(num_tokens, dim) + if self.use_l2: + nn.init.normal_(self.to_embedding.weight, std=1e-5) + else: + nn.init.kaiming_normal_(self.to_embedding.weight) + + def forward(self, x: Tensor) -> Tensor: + embedding = self.to_embedding(x) + return l2_norm(embedding) if self.use_l2 else embedding -- cgit v1.2.3-70-g09d2