diff options
Diffstat (limited to 'text_recognizer/networks/transformer')
4 files changed, 86 insertions, 78 deletions
diff --git a/text_recognizer/networks/transformer/attention.py b/text_recognizer/networks/transformer/attention.py index fca260d..85f513e 100644 --- a/text_recognizer/networks/transformer/attention.py +++ b/text_recognizer/networks/transformer/attention.py @@ -1,5 +1,5 @@ """Implementes the attention module for the transformer.""" -from typing import Optional, Tuple +from typing import Optional import torch import torch.nn.functional as F @@ -8,7 +8,6 @@ from torch import Tensor, einsum, nn from text_recognizer.networks.transformer.embeddings.rotary import ( RotaryEmbedding, - rotate_half, ) @@ -22,100 +21,75 @@ class Attention(nn.Module): causal: bool = False, dim_head: int = 64, dropout_rate: float = 0.0, - rotary_embedding: Optional[RotaryEmbedding] = None, ) -> None: super().__init__() - self.dim = dim + self.scale = self.dim**-0.5 self.num_heads = num_heads - self.causal = causal self.dim_head = dim_head + + self.causal = causal self.dropout_rate = dropout_rate - self.rotary_embedding = rotary_embedding - self.scale = self.dim**-0.5 - inner_dim = self.num_heads * self.dim_head + # Single key/value head + k_dim = dim_head + v_dim = dim_head + + out_dim = self.num_heads * self.dim_head - self.to_q = nn.Linear(self.dim, inner_dim, bias=False) - self.to_k = nn.Linear(self.dim, inner_dim, bias=False) - self.to_v = nn.Linear(self.dim, inner_dim, bias=False) + self.to_q = nn.Linear(self.dim, out_dim, bias=False) + self.to_k = nn.Linear(self.dim, k_dim, bias=False) + self.to_v = nn.Linear(self.dim, v_dim, bias=False) self.dropout = nn.Dropout(p=self.dropout_rate) # Feedforward - self.fc = nn.Linear(inner_dim, self.dim) + self.fc = nn.Linear(out_dim, self.dim) def forward( self, x: Tensor, context: Optional[Tensor] = None, - input_mask: Optional[Tensor] = None, - context_mask: Optional[Tensor] = None, + mask: Optional[Tensor] = None, + rotary_embedding: Optional[RotaryEmbedding] = None, ) -> Tensor: """Computes the attention.""" - b, n, _, device = *x.shape, x.device + b, device = x.shape[0], x.device q = self.to_q(x) + q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads) k = self.to_k(context) if context is not None else self.to_k(x) v = self.to_v(context) if context is not None else self.to_v(x) - q, k, v = map( - lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.num_heads), (q, k, v) - ) - if self.rotary_embedding is not None: - embedding = self.rotary_embedding(q) - q, k, v = _apply_rotary_emb(q, k, v, embedding[None, ...]) + if rotary_embedding is not None: + q, k, v = map(lambda t: rotary_embedding.rotate(t), (q, k, v)) - energy = einsum("b h i d, b h j d -> b h i j", q, k) * self.scale + energy = einsum("b h i d, b j d -> b h i j", q, k) * self.scale mask_value = -torch.finfo(energy.dtype).max - energy = apply_input_mask( - b, n, k, energy, input_mask, context, context_mask, mask_value, device - ) + energy = apply_input_mask(b, k, energy, mask, mask_value, device) if self.causal: - energy = apply_causal_mask(energy, input_mask, mask_value, device) + energy = apply_causal_mask(energy, mask, mask_value, device) attn = F.softmax(energy, dim=-1) attn = self.dropout(attn) - out = einsum("b h i j, b h j d -> b h i d", attn, v) + out = einsum("b h i j, b j d -> b h i d", attn, v) out = rearrange(out, "b h n d -> b n (h d)") out = self.fc(out) return out -def _apply_rotary_emb( - q: Tensor, k: Tensor, v: Tensor, freqs: Tensor -) -> Tuple[Tensor, Tensor, Tensor]: - q, k, v = map( - lambda t: (t * freqs.cos()) + (rotate_half(t) * freqs.sin()), (q, k, v) - ) - return q, k, v - - def apply_input_mask( b: int, - n: int, k: Tensor, energy: Tensor, - input_mask: Optional[Tensor], - context: Optional[Tensor], - context_mask: Optional[Tensor], + mask: Optional[Tensor], mask_value: Tensor, device: str, ) -> Tensor: """Applies an input mask.""" - if any(x is not None for x in (input_mask, context_mask)): - q_mask = ( - input_mask - if input_mask is not None - else torch.ones((b, n), device=device).bool() - ) - k_mask = q_mask if context is None else context_mask - k_mask = ( - torch.ones((b, k.shape[-2]), device=device).bool() - if k_mask is None - else k_mask - ) - q_mask = rearrange(q_mask, "b i -> b () i ()") + if mask is not None: + k_mask = torch.ones((b, k.shape[-2]), device=device).bool() + q_mask = rearrange(mask, "b i -> b () i ()") k_mask = rearrange(k_mask, "b j -> b () () j") input_mask = q_mask * k_mask diff --git a/text_recognizer/networks/transformer/decoder.py b/text_recognizer/networks/transformer/decoder.py index c7da226..741f5b3 100644 --- a/text_recognizer/networks/transformer/decoder.py +++ b/text_recognizer/networks/transformer/decoder.py @@ -6,16 +6,23 @@ from torch import Tensor, nn from text_recognizer.networks.transformer.attention import Attention from text_recognizer.networks.transformer.decoder_block import DecoderBlock +from text_recognizer.networks.transformer.embeddings.rotary import RotaryEmbedding from text_recognizer.networks.transformer.ff import FeedForward class Decoder(nn.Module): """Decoder Network.""" - def __init__(self, depth: int, dim: int, block: DecoderBlock) -> None: + def __init__( + self, + depth: int, + dim: int, + block: DecoderBlock, + rotary_embedding: Optional[RotaryEmbedding] = None, + ) -> None: super().__init__() self.depth = depth - self.has_pos_emb = block.has_pos_emb + self.rotary_embedding = rotary_embedding self.blocks = nn.ModuleList([deepcopy(block) for _ in range(self.depth)]) self.ln = nn.LayerNorm(dim) @@ -23,12 +30,14 @@ class Decoder(nn.Module): self, x: Tensor, context: Optional[Tensor] = None, - input_mask: Optional[Tensor] = None, - context_mask: Optional[Tensor] = None, + mask: Optional[Tensor] = None, ) -> Tensor: """Applies the network to the signals.""" for block in self.blocks: x = block( - x=x, context=context, input_mask=input_mask, context_mask=context_mask + x=x, + context=context, + mask=mask, + rotary_embedding=self.rotary_embedding, ) return self.ln(x) diff --git a/text_recognizer/networks/transformer/decoder_block.py b/text_recognizer/networks/transformer/decoder_block.py index b80a078..2dc4ddf 100644 --- a/text_recognizer/networks/transformer/decoder_block.py +++ b/text_recognizer/networks/transformer/decoder_block.py @@ -5,6 +5,7 @@ from typing import Optional, Type from torch import Tensor, nn from text_recognizer.networks.transformer.attention import Attention +from text_recognizer.networks.transformer.embeddings.rotary import RotaryEmbedding from text_recognizer.networks.transformer.ff import FeedForward @@ -25,22 +26,19 @@ class DecoderBlock(nn.Module): self.cross_attn = cross_attn self.ln_ff = deepcopy(norm) self.ff = ff - self.has_pos_emb = self.attn.rotary_embedding is not None def forward( self, x: Tensor, context: Optional[Tensor] = None, - input_mask: Optional[Tensor] = None, - context_mask: Optional[Tensor] = None, + mask: Optional[Tensor] = None, + rotary_embedding: Optional[RotaryEmbedding] = None, ) -> Tensor: """Applies decoder block on input signals.""" - x = x + self.attn(self.ln_attn(x), input_mask=input_mask) + x = x + self.attn(self.ln_attn(x), mask=mask, rotary_embedding=rotary_embedding) x = x + self.cross_attn( x=self.ln_cross_attn(x), context=context, - input_mask=input_mask, - context_mask=context_mask, ) x = x + self.ff(self.ln_ff(x)) return x diff --git a/text_recognizer/networks/transformer/embeddings/rotary.py b/text_recognizer/networks/transformer/embeddings/rotary.py index cc91206..ca0a260 100644 --- a/text_recognizer/networks/transformer/embeddings/rotary.py +++ b/text_recognizer/networks/transformer/embeddings/rotary.py @@ -6,6 +6,9 @@ Stolen from lucidrains: 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 @@ -17,24 +20,48 @@ class RotaryEmbedding(nn.Module): 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, x: Tensor) -> Tensor: + def forward(self, t: Tensor, cache_key: int) -> Tensor: """Encodes tensor x with rotary embeddings.""" - n = x.shape[-2] - t = torch.arange(n, device=x.device).type_as(self.inv_freqs) - freqs = torch.einsum("i , j -> i j", t, self.inv_freqs) - emb = torch.cat((freqs, freqs), dim=-1) - return emb[None, :, :] + 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: - if len(x.shape) == 3: - x = x.reshape((x.shape[0], -1, 2, x.shape[-1] // 2)) - else: - x = x.reshape((x.shape[0], x.shape[1], -1, 2, x.shape[-1] // 2)) - x1, x2 = x.unbind(dim=-2) - return torch.cat((-x2, x1), dim=-1) + 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_pos_emb(t: Tensor, freqs: Tensor) -> Tensor: - return (t * freqs.cos()) + (rotate_half(t) * freqs.sin()) +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) |