From 49ca6ade1a19f7f9c702171537fe4be0dfcda66d Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Fri, 25 Aug 2023 23:19:14 +0200
Subject: Rename and add flash atten
---
.../networks/transformer/embeddings/__init__.py | 1 -
.../networks/transformer/embeddings/axial.py | 104 ---------------------
.../networks/transformer/embeddings/rotary.py | 67 -------------
3 files changed, 172 deletions(-)
delete mode 100644 text_recognizer/networks/transformer/embeddings/__init__.py
delete mode 100644 text_recognizer/networks/transformer/embeddings/axial.py
delete mode 100644 text_recognizer/networks/transformer/embeddings/rotary.py
(limited to 'text_recognizer/networks/transformer/embeddings')
diff --git a/text_recognizer/networks/transformer/embeddings/__init__.py b/text_recognizer/networks/transformer/embeddings/__init__.py
deleted file mode 100644
index bb3f904..0000000
--- a/text_recognizer/networks/transformer/embeddings/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-"""Positional encodings for transformers."""
diff --git a/text_recognizer/networks/transformer/embeddings/axial.py b/text_recognizer/networks/transformer/embeddings/axial.py
deleted file mode 100644
index 9b872a9..0000000
--- a/text_recognizer/networks/transformer/embeddings/axial.py
+++ /dev/null
@@ -1,104 +0,0 @@
-"""Axial attention for multi-dimensional data.
-
-Stolen from:
- https://github.com/lucidrains/axial-attention/blob/
- eff2c10c2e76c735a70a6b995b571213adffbbb7/axial_attention/axial_attention.py#L100
-"""
-from functools import reduce
-from operator import mul
-from typing import Optional, Sequence
-
-import torch
-from torch import Tensor, nn
-
-
-class AxialPositionalEmbedding(nn.Module):
- def __init__(
- self,
- dim: int,
- axial_shape: Sequence[int],
- axial_dims: Optional[Sequence[int]] = None,
- ) -> None:
- super().__init__()
-
- self.dim = dim
- self.shape = axial_shape
- self.max_seq_len = reduce(mul, axial_shape, 1)
-
- self.summed = axial_dims is None
- axial_dims = ((dim,) * len(axial_shape)) if self.summed else axial_dims
-
- assert len(self.shape) == len(
- axial_dims
- ), "number of axial dimensions must equal the number of dimensions in the shape"
- assert (
- self.summed or not self.summed and sum(axial_dims) == dim
- ), f"axial dimensions must sum up to the target dimension {dim}"
-
- self.weights = ParameterList(self, "weights", len(axial_shape))
-
- for ind, (shape, axial_dim) in enumerate(zip(self.shape, axial_dims)):
- ax_shape = [1] * len(self.shape)
- ax_shape[ind] = shape
- ax_shape = (1, *ax_shape, axial_dim)
- ax_emb = nn.Parameter(torch.zeros(ax_shape).normal_(0, 1))
- self.weights.append(ax_emb)
-
- def forward(self, x: Tensor) -> Tensor:
- """Returns axial positional embedding."""
- b, t, _ = x.shape
- assert (
- t <= self.max_seq_len
- ), f"Sequence length ({t}) must be less than the maximum sequence length allowed ({self.max_seq_len})"
- embs = []
-
- for ax_emb in self.weights.to_list():
- axial_dim = ax_emb.shape[-1]
- expand_shape = (b, *self.shape, axial_dim)
- emb = ax_emb.expand(expand_shape).reshape(b, self.max_seq_len, axial_dim)
- embs.append(emb)
-
- pos_emb = sum(embs) if self.summed else torch.cat(embs, dim=-1)
- return pos_emb[:, :t].to(x)
-
-
-# a mock parameter list object until below issue is resolved
-# https://github.com/pytorch/pytorch/issues/36035
-class ParameterList(object):
- def __init__(self, kls, prefix, length):
- self.ind = 0
- self.kls = kls
- self.prefix = prefix
- self.length = length
-
- def _keyname(self, prefix, ind):
- return f"{prefix}_{ind}"
-
- def append(self, x):
- setattr(self.kls, self._keyname(self.prefix, self.ind), x)
- self.ind += 1
-
- def to_list(self):
- return [
- getattr(self.kls, self._keyname(self.prefix, i)) for i in range(self.length)
- ]
-
-
-class AxialPositionalEmbeddingImage(nn.Module):
- def __init__(
- self,
- dim: int,
- axial_shape: Sequence[int],
- axial_dims: Optional[Sequence[int]] = None,
- ) -> None:
- super().__init__()
- axial_dims = (dim // 2, dim // 2) if axial_dims is None else axial_dims
- assert len(axial_shape) == 2, "Axial shape must have 2 dimensions for images"
- self.dim = dim
- self.pos_emb = AxialPositionalEmbedding(dim, axial_shape, axial_dims)
-
- def forward(self, img):
- b, c, h, w = img.shape
- img = img.permute(0, 2, 3, 1).reshape(b, h * w, c)
- pos_emb = self.pos_emb(img)
- return pos_emb.reshape(b, h, w, self.dim).permute(0, 3, 1, 2)
diff --git a/text_recognizer/networks/transformer/embeddings/rotary.py b/text_recognizer/networks/transformer/embeddings/rotary.py
deleted file mode 100644
index ca0a260..0000000
--- a/text_recognizer/networks/transformer/embeddings/rotary.py
+++ /dev/null
@@ -1,67 +0,0 @@
-"""Roatary embedding.
-
-Stolen from lucidrains:
- https://github.com/lucidrains/rotary-embedding-torch
-
-Explanation of roatary:
- https://blog.eleuther.ai/rotary-embeddings/
-"""
-from inspect import isfunction
-
-from einops import rearrange, repeat
-import torch
-from torch import Tensor, nn
-
-
-class RotaryEmbedding(nn.Module):
- """Rotary positional embedding."""
-
- def __init__(self, dim: int) -> None:
- super().__init__()
- inv_freqs = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
- self.register_buffer("inv_freqs", inv_freqs)
- self.cache = {}
-
- def rotate(self, t: Tensor, dim: int = -2) -> Tensor:
- """Rotate vector."""
- device, n = t.device, t.shape[dim]
- freqs = self.forward(lambda: torch.arange(n, device=device), cache_key=n)
- return apply_rotary_emb(t, freqs)
-
- def forward(self, t: Tensor, cache_key: int) -> Tensor:
- """Encodes tensor x with rotary embeddings."""
- if cache_key in self.cache:
- return self.cache[cache_key]
-
- if isfunction(t):
- t = t()
-
- freqs = self.inv_freqs
- freqs = torch.einsum("..., f -> ... f", t.type(freqs.dtype), freqs)
- freqs = repeat(freqs, "... n -> ... (n r)", r=2)
- self.cache[cache_key] = freqs
- return freqs
-
-
-def rotate_half(x: Tensor) -> Tensor:
- x = rearrange(x, "... (d r) -> ... d r", r=2)
- x1, x2 = x.unbind(dim=-1)
- x = torch.stack((-x2, x1), dim=-1)
- return rearrange(x, "... d r -> ... (d r)")
-
-
-def apply_rotary_emb(t: Tensor, freqs: Tensor, start_index: int = 0) -> Tensor:
- freqs = freqs.to(t)
- rot_dim = freqs.shape[-1]
- end_index = start_index + rot_dim
- assert rot_dim <= t.shape[-1], (
- f"feature dimension {t.shape[-1]} is not of sufficient size to rotate"
- f"in all the positions {rot_dim}"
- )
- t_left, t, t_right = (
- t[..., :start_index],
- t[..., start_index:end_index],
- t[..., end_index:],
- )
- t = (t * freqs.cos()) + (rotate_half(t) * freqs.sin())
- return torch.cat((t_left, t, t_right), dim=-1)
--
cgit v1.2.3-70-g09d2