diff options
Diffstat (limited to 'text_recognizer/networks/transformer/embeddings')
3 files changed, 0 insertions, 172 deletions
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) |