diff options
| -rw-r--r-- | notebooks/00-testing-stuff-out.ipynb | 163 | ||||
| -rw-r--r-- | text_recognizer/models/vqvae.py | 12 | ||||
| -rw-r--r-- | text_recognizer/networks/transformer/attention.py | 72 | ||||
| -rw-r--r-- | training/conf/network/vqvae.yaml | 6 |
4 files changed, 233 insertions, 20 deletions
diff --git a/notebooks/00-testing-stuff-out.ipynb b/notebooks/00-testing-stuff-out.ipynb index 7c7b3a6..92faaf7 100644 --- a/notebooks/00-testing-stuff-out.ipynb +++ b/notebooks/00-testing-stuff-out.ipynb @@ -358,6 +358,169 @@ "metadata": {}, "outputs": [], "source": [] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [], + "source": [ + "b, n = 16, 128\n", + "device = \"cpu\"" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [], + "source": [ + "x = lambda: torch.ones((b, n), device=device).bool()" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "torch.Size([16, 128])" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "x().shape" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "torch.Size([16, 128])" + ] + }, + "execution_count": 12, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "torch.ones((b, n), device=device).bool().shape" + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "metadata": {}, + "outputs": [], + "source": [ + "x = torch.randn(1, 1, 576, 640)" + ] + }, + { + "cell_type": "code", + "execution_count": 30, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "144" + ] + }, + "execution_count": 30, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "576 // 4" + ] + }, + { + "cell_type": "code", + "execution_count": 29, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "160" + ] + }, + "execution_count": 29, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "640 // 4" + ] + }, + { + "cell_type": "code", + "execution_count": 31, + "metadata": {}, + "outputs": [], + "source": [ + "x = torch.randn(1, 1, 144, 160)" + ] + }, + { + "cell_type": "code", + "execution_count": 32, + "metadata": {}, + "outputs": [], + "source": [ + "from einops import rearrange" + ] + }, + { + "cell_type": "code", + "execution_count": 35, + "metadata": {}, + "outputs": [], + "source": [ + "patch_size=4\n", + "p = rearrange(x, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_size, p2 = patch_size)" + ] + }, + { + "cell_type": "code", + "execution_count": 36, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "torch.Size([1, 1440, 16])" + ] + }, + "execution_count": 36, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "p.shape" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { diff --git a/text_recognizer/models/vqvae.py b/text_recognizer/models/vqvae.py index ef2213c..078235e 100644 --- a/text_recognizer/models/vqvae.py +++ b/text_recognizer/models/vqvae.py @@ -29,12 +29,14 @@ class LitVQVAEModel(LitBaseModel): """Forward pass with the transformer network.""" return self.network.predict(data) - def _log_prediction(self, data: Tensor, reconstructions: Tensor) -> None: + def _log_prediction( + self, data: Tensor, reconstructions: Tensor, title: str + ) -> None: """Logs prediction on image with wandb.""" try: self.logger.experiment.log( { - "val_pred_examples": [ + title: [ wandb.Image(data[0]), wandb.Image(reconstructions[0]), ] @@ -59,7 +61,8 @@ class LitVQVAEModel(LitBaseModel): loss = self.loss_fn(reconstructions, data) loss += vq_loss self.log("val_loss", loss, prog_bar=True) - self._log_prediction(data, reconstructions) + title = "val_pred_examples" + self._log_prediction(data, reconstructions, title) def test_step(self, batch: Tuple[Tensor, Tensor], batch_idx: int) -> None: """Test step.""" @@ -67,4 +70,5 @@ class LitVQVAEModel(LitBaseModel): reconstructions, vq_loss = self.network(data) loss = self.loss_fn(reconstructions, data) loss += vq_loss - self._log_prediction(data, reconstructions) + title = "test_pred_examples" + self._log_prediction(data, reconstructions, title) diff --git a/text_recognizer/networks/transformer/attention.py b/text_recognizer/networks/transformer/attention.py index 8724691..623d680 100644 --- a/text_recognizer/networks/transformer/attention.py +++ b/text_recognizer/networks/transformer/attention.py @@ -1,9 +1,10 @@ """Implementes the attention module for the transformer.""" from typing import Optional, Tuple +from einops import rearrange from einops.layers.torch import Rearrange -import numpy as np import torch +from torch import einsum from torch import nn from torch import Tensor import torch.nn.functional as F @@ -34,7 +35,7 @@ class Attention(nn.Module): self.attn_fn = F.softmax # Feedforward - self.proj = nn.Linear(inner_dim, dim) + self.fc = nn.Linear(inner_dim, dim) @staticmethod def _apply_rotary_emb( @@ -47,8 +48,42 @@ class Attention(nn.Module): k = torch.cat((kl, kr), dim=-1) return q, k - def _cross_attention(self) -> Tensor: - pass + @staticmethod + def _compute_input_mask( + b: int, + n: int, + k: Tensor, + mask: Optional[Tensor], + context: Optional[Tensor], + context_mask: Optional[Tensor], + device: str, + ) -> Optional[Tensor]: + if any(x is not None for x in (mask, context_mask)): + q_mask = ( + mask if mask is not None else torch.ones((b, n), device=device).bool() + ) + k_mask = q_mask if context is not 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 ()") + k_mask = rearrange(k_mask, "b i -> b () () j") + return q_mask * k_mask + return + + @staticmethod + def _apply_causal_mask( + energy: Tensor, mask: Tensor, mask_value: Tensor, device: str + ) -> Tensor: + i, j = energy.shape[-2:] + r = torch.arange(i, device=device) + mask = rearrange(r, "i -> () () i ()") < rearrange(r, "j -> () () () j") + mask = F.pad(mask, (j - i, 0), value=False) + energy.masked_fill_(mask, mask_value) + del mask + return energy def forward( self, @@ -67,14 +102,25 @@ class Attention(nn.Module): k, ) - input_mask = None - if any(x is not None for x in (mask, context_mask)): - q_mask = ( - mask - if mask is not None - else lambda: torch.ones((b, n), device=device).bool() - ) - pass + input_mask = self._compute_input_mask( + b, n, k, mask, context, context_mask, device + ) # Compute the attention - energy = (q @ k.transpose(-2, -1)) * self.scale + energy = einsum("b h i d, b h j d -> b h i j", q, k) * self.scale + mask_value = -torch.finfo(energy.dtype).max + + # Apply input mask + if input_mask is not None: + energy = energy.masked_fill_(~input_mask, mask_value) + del input_mask + + if self.causal: + energy = self._apply_causal_mask(energy, mask, mask_value, device) + + attn = self.attn_fn(energy, dim=-1) + attn = self.dropout(attn) + out = einsum("b h i j, b h 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, attn diff --git a/training/conf/network/vqvae.yaml b/training/conf/network/vqvae.yaml index 8c30bbd..288d2aa 100644 --- a/training/conf/network/vqvae.yaml +++ b/training/conf/network/vqvae.yaml @@ -2,9 +2,9 @@ type: VQVAE args: in_channels: 1 - channels: [32, 64, 64] - kernel_sizes: [4, 4, 4] - strides: [2, 2, 2] + channels: [64, 96] + kernel_sizes: [4, 4] + strides: [2, 2] num_residual_layers: 2 embedding_dim: 64 num_embeddings: 256 |