From 58ae7154aa945cfe5a46592cc1dfb28f0a4e51b3 Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Sat, 1 May 2021 23:53:50 +0200 Subject: Working on new attention module --- text_recognizer/networks/transformer/norm.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) (limited to 'text_recognizer/networks/transformer/norm.py') diff --git a/text_recognizer/networks/transformer/norm.py b/text_recognizer/networks/transformer/norm.py index 99a5291..9160876 100644 --- a/text_recognizer/networks/transformer/norm.py +++ b/text_recognizer/networks/transformer/norm.py @@ -20,3 +20,16 @@ class Rezero(nn.Module): def forward(self, x: Tensor, **kwargs: Dict) -> Tensor: x, *rest = self.fn(x, **kwargs) return (x * self.g, *rest) + + +class ScaleNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1.0e-5) -> None: + super().__init__() + self.scale = dim ** -0.5 + self.eps = eps + self.g = nn.Parameter(torch.ones(1)) + + def forward(self, x: Tensor) -> Tensor: + norm = torch.norm(x, dim=-1, keepdim=True) * self.scale + return x / norm.clamp(min=self.eps) self.g + -- cgit v1.2.3-70-g09d2