diff options
15 files changed, 25 insertions, 29 deletions
diff --git a/text_recognizer/networks/conv_transformer.py b/text_recognizer/networks/conv_transformer.py index 40047ad..d2ab8ea 100644 --- a/text_recognizer/networks/conv_transformer.py +++ b/text_recognizer/networks/conv_transformer.py @@ -1,7 +1,7 @@ """Base network module.""" from typing import Optional, Tuple, Type -from torch import nn, Tensor +from torch import Tensor, nn from text_recognizer.networks.transformer.decoder import Decoder from text_recognizer.networks.transformer.embeddings.axial import ( diff --git a/text_recognizer/networks/convnext/__init__.py b/text_recognizer/networks/convnext/__init__.py index 1743f09..faebe6f 100644 --- a/text_recognizer/networks/convnext/__init__.py +++ b/text_recognizer/networks/convnext/__init__.py @@ -1,7 +1,7 @@ """Convnext module.""" -from text_recognizer.networks.convnext.convnext import ConvNext from text_recognizer.networks.convnext.attention import ( Attention, FeedForward, TransformerBlock, ) +from text_recognizer.networks.convnext.convnext import ConvNext diff --git a/text_recognizer/networks/convnext/attention.py b/text_recognizer/networks/convnext/attention.py index 7f03436..a373c37 100644 --- a/text_recognizer/networks/convnext/attention.py +++ b/text_recognizer/networks/convnext/attention.py @@ -1,8 +1,8 @@ """Convolution self attention block.""" -from einops import reduce, rearrange -from torch import einsum, nn, Tensor import torch.nn.functional as F +from einops import rearrange, reduce +from torch import Tensor, einsum, nn from text_recognizer.networks.convnext.norm import LayerNorm from text_recognizer.networks.convnext.residual import Residual diff --git a/text_recognizer/networks/convnext/downsample.py b/text_recognizer/networks/convnext/downsample.py index 6e4306f..c28ecca 100644 --- a/text_recognizer/networks/convnext/downsample.py +++ b/text_recognizer/networks/convnext/downsample.py @@ -1,7 +1,7 @@ from typing import Tuple from einops.layers.torch import Rearrange -from torch import nn, Tensor +from torch import Tensor, nn class Downsample(nn.Module): diff --git a/text_recognizer/networks/convnext/norm.py b/text_recognizer/networks/convnext/norm.py index 2d896e5..23cf07a 100644 --- a/text_recognizer/networks/convnext/norm.py +++ b/text_recognizer/networks/convnext/norm.py @@ -1,6 +1,6 @@ """Layer norm for conv layers.""" import torch -from torch import nn, Tensor +from torch import Tensor, nn class LayerNorm(nn.Module): diff --git a/text_recognizer/networks/convnext/residual.py b/text_recognizer/networks/convnext/residual.py index 3f44390..8e76ae9 100644 --- a/text_recognizer/networks/convnext/residual.py +++ b/text_recognizer/networks/convnext/residual.py @@ -1,6 +1,7 @@ """Generic residual layer.""" from typing import Callable -from torch import nn, Tensor + +from torch import Tensor, nn class Residual(nn.Module): 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)) |