feat(norm): add prenorm

1 files changed, 16 insertions, 0 deletions

diff --git a/text_recognizer/networks/transformer/norm.py b/text_recognizer/networks/transformer/norm.py
index 2b416e6..be38346 100644
--- a/text_recognizer/networks/transformer/norm.py
+++ b/text_recognizer/networks/transformer/norm.py
@@ -4,6 +4,8 @@ Copied from lucidrains:
     https://github.com/lucidrains/x-transformers/blob/main/x_transformers/x_transformers.py
 
 """
+from typing import Dict, Type
+
 import torch
 from torch import nn
 from torch import Tensor
@@ -22,3 +24,17 @@ class RMSNorm(nn.Module):
         """Applies normalization."""
         norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
         return x / norm.clamp(min=self.eps) * self.g
+
+
+class PreNorm(nn.Module):
+    """Applies layer normalization then function."""
+
+    def __init__(self, normalized_shape: int, fn: Type[nn.Module]) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(normalized_shape)
+        self.fn = fn
+
+    def forward(self, x: Tensor, **kwargs: Dict) -> Tensor:
+        """Applies pre norm."""
+        x = self.norm(x)
+        return self.fn(x, **kwargs)