From 7ae1f8f9654dcea0a9a22310ac0665a5d3202f0f Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Mon, 26 Apr 2021 22:04:47 +0200 Subject: Reformatting transformer (work in progress) --- text_recognizer/networks/transformer/mlp.py | 46 +++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) create mode 100644 text_recognizer/networks/transformer/mlp.py (limited to 'text_recognizer/networks/transformer/mlp.py') diff --git a/text_recognizer/networks/transformer/mlp.py b/text_recognizer/networks/transformer/mlp.py new file mode 100644 index 0000000..4028ab3 --- /dev/null +++ b/text_recognizer/networks/transformer/mlp.py @@ -0,0 +1,46 @@ +"""Feedforward layer in transformer. + +Stolen from lucidrains: + https://github.com/lucidrains/x-transformers/blob/main/x_transformers/x_transformers.py +""" +from typing import Optional + +from torch import nn +from torch import Tensor +import torch.nn.functional as F + + +class GEGLU(nn.Module): + def __init__(self, dim_in: int, dim_out: int) -> None: + super().__init__() + self.fc = nn.Linear(dim_in, dim_out * 2) + + def forward(self, x: Tensor) -> Tensor: + x, gate = self.fc(x).chunk(2, dim=-1) + return x * F.gelu(gate) + + +class FeedForward(nn.Module): + def __init__( + self, + dim: int, + dim_out: Optional[int] = None, + expansion_factor: int = 4, + glu: bool = True, + dropout_rate: float = 0.0, + ) -> None: + super().__init__() + inner_dim = dim * expansion_factor + dim_out = dim_out if dim_out is not None else dim + in_projection = ( + nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU()) + if not glu + else GEGLU(dim, inner_dim) + ) + + self.mlp = nn.Sequential( + in_projection, nn.Dropout(dropout_rate), nn.Linear(inner_dim, dim_out) + ) + + def forward(self, x: Tensor) -> Tensor: + return self.mlp(x) -- cgit v1.2.3-70-g09d2