summaryrefslogtreecommitdiff
path: root/text_recognizer/networks/transformer
diff options
context:
space:
mode:
authorGustaf Rydholm <gustaf.rydholm@gmail.com>2022-09-18 18:11:21 +0200
committerGustaf Rydholm <gustaf.rydholm@gmail.com>2022-09-18 18:11:53 +0200
commit2cc6aa059139b57057609817913ad515063c2eab (patch)
tree5433f69a5eaf63e064a100bf900783127c7b1ff4 /text_recognizer/networks/transformer
parent88caa5c466225d4752541c352c5777235f8f0c61 (diff)
Format imports
Format imports
Diffstat (limited to 'text_recognizer/networks/transformer')
-rw-r--r--text_recognizer/networks/transformer/__init__.py4
-rw-r--r--text_recognizer/networks/transformer/attention.py8
-rw-r--r--text_recognizer/networks/transformer/decoder.py2
-rw-r--r--text_recognizer/networks/transformer/embeddings/absolute.py6
-rw-r--r--text_recognizer/networks/transformer/embeddings/axial.py5
-rw-r--r--text_recognizer/networks/transformer/embeddings/fourier.py3
-rw-r--r--text_recognizer/networks/transformer/embeddings/rotary.py3
-rw-r--r--text_recognizer/networks/transformer/ff.py3
-rw-r--r--text_recognizer/networks/transformer/norm.py5
9 files changed, 17 insertions, 22 deletions
diff --git a/text_recognizer/networks/transformer/__init__.py b/text_recognizer/networks/transformer/__init__.py
index af08aa8..0d17deb 100644
--- a/text_recognizer/networks/transformer/__init__.py
+++ b/text_recognizer/networks/transformer/__init__.py
@@ -1,6 +1,6 @@
"""Transformer modules."""
from text_recognizer.networks.transformer.attention import Attention
-from text_recognizer.networks.transformer.ff import FeedForward
-from text_recognizer.networks.transformer.norm import RMSNorm
from text_recognizer.networks.transformer.decoder import Decoder, DecoderBlock
from text_recognizer.networks.transformer.embeddings.rotary import RotaryEmbedding
+from text_recognizer.networks.transformer.ff import FeedForward
+from text_recognizer.networks.transformer.norm import RMSNorm
diff --git a/text_recognizer/networks/transformer/attention.py b/text_recognizer/networks/transformer/attention.py
index 3df5333..fca260d 100644
--- a/text_recognizer/networks/transformer/attention.py
+++ b/text_recognizer/networks/transformer/attention.py
@@ -1,12 +1,10 @@
"""Implementes the attention module for the transformer."""
from typing import Optional, Tuple
-from einops import rearrange
import torch
-from torch import einsum
-from torch import nn
-from torch import Tensor
import torch.nn.functional as F
+from einops import rearrange
+from torch import Tensor, einsum, nn
from text_recognizer.networks.transformer.embeddings.rotary import (
RotaryEmbedding,
@@ -35,7 +33,7 @@ class Attention(nn.Module):
self.dropout_rate = dropout_rate
self.rotary_embedding = rotary_embedding
- self.scale = self.dim ** -0.5
+ self.scale = self.dim**-0.5
inner_dim = self.num_heads * self.dim_head
self.to_q = nn.Linear(self.dim, inner_dim, bias=False)
diff --git a/text_recognizer/networks/transformer/decoder.py b/text_recognizer/networks/transformer/decoder.py
index 8ea021b..efa1e89 100644
--- a/text_recognizer/networks/transformer/decoder.py
+++ b/text_recognizer/networks/transformer/decoder.py
@@ -2,7 +2,7 @@
from copy import deepcopy
from typing import Optional, Type
-from torch import nn, Tensor
+from torch import Tensor, nn
from text_recognizer.networks.transformer.attention import Attention
from text_recognizer.networks.transformer.ff import FeedForward
diff --git a/text_recognizer/networks/transformer/embeddings/absolute.py b/text_recognizer/networks/transformer/embeddings/absolute.py
index e5cdc18..9274b55 100644
--- a/text_recognizer/networks/transformer/embeddings/absolute.py
+++ b/text_recognizer/networks/transformer/embeddings/absolute.py
@@ -1,9 +1,9 @@
"""Absolute positional embedding."""
-from einops import rearrange
import torch
-from torch import nn
import torch.nn.functional as F
+from einops import rearrange
+from torch import nn
def l2norm(t, groups=1):
@@ -15,7 +15,7 @@ def l2norm(t, groups=1):
class AbsolutePositionalEmbedding(nn.Module):
def __init__(self, dim, max_seq_len, l2norm_embed=False):
super().__init__()
- self.scale = dim ** -0.5 if not l2norm_embed else 1.0
+ self.scale = dim**-0.5 if not l2norm_embed else 1.0
self.max_seq_len = max_seq_len
self.l2norm_embed = l2norm_embed
self.emb = nn.Embedding(max_seq_len, dim)
diff --git a/text_recognizer/networks/transformer/embeddings/axial.py b/text_recognizer/networks/transformer/embeddings/axial.py
index 7b84e12..25d8f60 100644
--- a/text_recognizer/networks/transformer/embeddings/axial.py
+++ b/text_recognizer/networks/transformer/embeddings/axial.py
@@ -3,10 +3,11 @@
Stolen from:
https://github.com/lucidrains/axial-attention/blob/eff2c10c2e76c735a70a6b995b571213adffbbb7/axial_attention/axial_attention.py#L100
"""
+from functools import reduce
+from operator import mul
+
import torch
from torch import nn
-from operator import mul
-from functools import reduce
class AxialPositionalEmbedding(nn.Module):
diff --git a/text_recognizer/networks/transformer/embeddings/fourier.py b/text_recognizer/networks/transformer/embeddings/fourier.py
index 7843c60..28da7a1 100644
--- a/text_recognizer/networks/transformer/embeddings/fourier.py
+++ b/text_recognizer/networks/transformer/embeddings/fourier.py
@@ -1,8 +1,7 @@
"""Fourier positional embedding."""
import numpy as np
import torch
-from torch import nn
-from torch import Tensor
+from torch import Tensor, nn
class PositionalEncoding(nn.Module):
diff --git a/text_recognizer/networks/transformer/embeddings/rotary.py b/text_recognizer/networks/transformer/embeddings/rotary.py
index 722478e..cc91206 100644
--- a/text_recognizer/networks/transformer/embeddings/rotary.py
+++ b/text_recognizer/networks/transformer/embeddings/rotary.py
@@ -7,8 +7,7 @@ Explanation of roatary:
https://blog.eleuther.ai/rotary-embeddings/
"""
import torch
-from torch import nn
-from torch import Tensor
+from torch import Tensor, nn
class RotaryEmbedding(nn.Module):
diff --git a/text_recognizer/networks/transformer/ff.py b/text_recognizer/networks/transformer/ff.py
index 4028ab3..3ccf5b5 100644
--- a/text_recognizer/networks/transformer/ff.py
+++ b/text_recognizer/networks/transformer/ff.py
@@ -5,9 +5,8 @@ Stolen from lucidrains:
"""
from typing import Optional
-from torch import nn
-from torch import Tensor
import torch.nn.functional as F
+from torch import Tensor, nn
class GEGLU(nn.Module):
diff --git a/text_recognizer/networks/transformer/norm.py b/text_recognizer/networks/transformer/norm.py
index 4cd3b5b..0bd2e16 100644
--- a/text_recognizer/networks/transformer/norm.py
+++ b/text_recognizer/networks/transformer/norm.py
@@ -7,8 +7,7 @@ Copied from lucidrains:
from typing import Dict, Optional, Type
import torch
-from torch import nn
-from torch import Tensor
+from torch import Tensor, nn
class RMSNorm(nn.Module):
@@ -16,7 +15,7 @@ class RMSNorm(nn.Module):
def __init__(self, dim: int, eps: float = 1e-8) -> None:
super().__init__()
- self.scale = dim ** -0.5
+ self.scale = dim**-0.5
self.eps = eps
self.g = nn.Parameter(torch.ones(dim))