Reformatting transformer (work in progress)

1 files changed, 22 insertions, 0 deletions

diff --git a/text_recognizer/networks/transformer/norm.py b/text_recognizer/networks/transformer/norm.py
new file mode 100644
index 0000000..99a5291
--- /dev/null
+++ b/text_recognizer/networks/transformer/norm.py
@@ -0,0 +1,22 @@
+"""Normalization layers for transfromers.
+
+Copied from lucidrains:
+    https://github.com/lucidrains/x-transformers/blob/main/x_transformers/x_transformers.py
+
+"""
+from typing import Callable, Dict
+
+import torch
+from torch import nn
+from torch import Tensor
+
+
+class Rezero(nn.Module):
+    def __init__(self, fn: Callable) -> None:
+        super().__init__()
+        self.fn = fn
+        self.g = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x: Tensor, **kwargs: Dict) -> Tensor:
+        x, *rest = self.fn(x, **kwargs)
+        return (x * self.g, *rest)