Complete cnn-transformer network, not tested

5 files changed, 161 insertions, 346 deletions

diff --git a/notebooks/00-scratch-pad.ipynb b/notebooks/00-scratch-pad.ipynb
index 1e30038..2c98064 100644
--- a/notebooks/00-scratch-pad.ipynb
+++ b/notebooks/00-scratch-pad.ipynb
@@ -2,9 +2,18 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 5,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "The autoreload extension is already loaded. To reload it, use:\n",
+      "  %reload_ext autoreload\n"
+     ]
+    }
+   ],
    "source": [
     "%load_ext autoreload\n",
     "%autoreload 2\n",
@@ -30,472 +39,230 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from pathlib import Path"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
-    "import attr"
+    "from text_recognizer.networks.encoders.efficientnet.efficientnet import EfficientNet"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 44,
-   "metadata": {},
+   "execution_count": 7,
+   "metadata": {
+    "scrolled": false
+   },
    "outputs": [],
    "source": [
-    "@attr.s\n",
-    "class B(nn.Module):\n",
-    "    input_dim = attr.ib()\n",
-    "    hidden = attr.ib()\n",
-    "    xx = attr.ib(init=False, default=\"hek\")\n",
-    "    \n",
-    "    def __attrs_post_init__(self):\n",
-    "        super().__init__()\n",
-    "        self.fc = nn.Linear(self.input_dim, self.hidden)\n",
-    "        self.xx = \"da\"\n",
-    "    \n",
-    "    def forward(self, x):\n",
-    "        return self.fc(x)"
+    "en = EfficientNet(\"b0\")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 49,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [],
    "source": [
-    "def f(x):\n",
-    "    return 2\n",
-    "\n",
-    "@attr.s(auto_attribs=True)\n",
-    "class T(B):\n",
-    "    \n",
-    "    h: Path = attr.ib(converter=Path)\n",
-    "    p: int = attr.ib(init=False, default=f(3))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 53,
-   "metadata": {},
-   "outputs": [
-    {
-     "ename": "TypeError",
-     "evalue": "__init__() missing 1 required positional argument: 'hidden'",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
-      "\u001b[0;32m<ipython-input-53-ef8b390156f4>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mT\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput_dim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m16\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mh\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"hej\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
-      "\u001b[0;31mTypeError\u001b[0m: __init__() missing 1 required positional argument: 'hidden'"
-     ]
-    }
-   ],
-   "source": [
-    "t = T(input_dim=16, h=\"hej\")"
+    "def generate_square_subsequent_mask(size: int) -> torch.Tensor:\n",
+    "    \"\"\"Generate a triangular (size, size) mask.\"\"\"\n",
+    "    mask = (torch.triu(torch.ones(size, size)) == 1).transpose(0, 1)\n",
+    "    mask = mask.float().masked_fill(mask == 0, float(\"-inf\")).masked_fill(mask == 1, float(0.0))\n",
+    "    return mask"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 51,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "'da'"
+       "tensor([[0., -inf, -inf, -inf],\n",
+       "        [0., 0., -inf, -inf],\n",
+       "        [0., 0., 0., -inf],\n",
+       "        [0., 0., 0., 0.]])"
       ]
      },
-     "execution_count": 51,
+     "execution_count": 10,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "t.xx"
+    "generate_square_subsequent_mask(4)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from torch import Tensor"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 52,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "2"
-      ]
-     },
-     "execution_count": 52,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
-    "t.p"
+    "tgt = torch.randint(0, 4, (1, 4))\n",
+    "tgt_mask = torch.ones_like(tgt).bool()"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 53,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "16"
+       "tensor([[True, True, True, True]])"
       ]
      },
-     "execution_count": 19,
+     "execution_count": 53,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "t.input_dim"
+    "tgt_mask"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 13,
    "metadata": {},
    "outputs": [],
    "source": [
-    "x = torch.rand(16, 16)"
+    "def target_padding_mask(trg: Tensor, pad_index: int) -> Tensor:\n",
+    "    \"\"\"Returns causal target mask.\"\"\"\n",
+    "    trg_pad_mask = (trg != pad_index)[:, None, None]\n",
+    "    trg_len = trg.shape[1]\n",
+    "    trg_sub_mask = torch.tril(torch.ones((trg_len, trg_len), device=trg.device)).bool()\n",
+    "    trg_mask = trg_pad_mask & trg_sub_mask\n",
+    "    return trg_mask"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": 54,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "torch.Size([16, 16])"
-      ]
-     },
-     "execution_count": 21,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
-    "x.shape"
+    "t = torch.randint(0, 6, (0, 4))"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 23,
+   "execution_count": 55,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "T(input_dim=16, hidden=24, h=PosixPath('hej'))"
-      ]
-     },
-     "execution_count": 23,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
-    "t.cuda()"
+    "t = torch.Tensor([[0, 0, 0, 3, 3, 3]])"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 24,
+   "execution_count": 58,
    "metadata": {},
    "outputs": [],
    "source": [
-    "x = x.cuda()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 25,
-   "metadata": {
-    "scrolled": true
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "tensor([[ 3.6047e-01,  1.0200e+00,  3.6786e-01,  1.6077e-01,  3.9281e-02,\n",
-       "          3.2830e-01,  1.3433e-01, -9.0334e-02, -3.8712e-01,  8.1547e-01,\n",
-       "         -5.4483e-01, -9.7471e-01,  3.3706e-01, -9.5283e-01, -1.6271e-01,\n",
-       "          3.8504e-01, -5.0106e-01, -4.8638e-01,  3.7033e-01, -4.9557e-01,\n",
-       "          2.6555e-01,  5.1245e-01,  6.6751e-01, -2.6291e-01],\n",
-       "        [ 1.3811e-01,  7.4522e-01,  4.9935e-01,  3.3878e-01,  1.8501e-01,\n",
-       "          2.2269e-02, -2.0328e-01,  1.4629e-01, -2.2957e-01,  4.1197e-01,\n",
-       "         -1.9555e-01, -4.7609e-01,  9.0206e-02, -8.8568e-01, -2.1618e-01,\n",
-       "          2.8882e-01, -5.4335e-01, -6.6301e-01,  4.9990e-01, -4.0144e-01,\n",
-       "          3.6403e-01,  5.3901e-01,  8.6665e-01, -7.8312e-02],\n",
-       "        [ 1.6493e-02,  4.6157e-01,  2.9500e-02,  2.4190e-01,  6.5753e-01,\n",
-       "          4.3770e-02, -5.3773e-02,  1.8183e-01, -2.5983e-02,  4.1634e-01,\n",
-       "         -3.5218e-01, -5.6129e-01,  4.1452e-01, -1.2265e+00, -5.8544e-01,\n",
-       "          3.6382e-01, -6.4090e-01, -5.8679e-01,  4.3489e-02, -1.1233e-01,\n",
-       "          3.1175e-01,  4.2857e-01,  1.6501e-01, -2.4118e-01],\n",
-       "        [ 9.2361e-02,  6.0196e-01,  1.3081e-02, -8.1091e-02,  4.2342e-01,\n",
-       "         -8.8457e-02, -8.1851e-02, -1.1562e-01, -1.5049e-01,  4.9972e-01,\n",
-       "         -3.0432e-01, -7.8619e-01,  2.1060e-01, -1.0598e+00, -4.6542e-01,\n",
-       "          4.2382e-01, -6.5671e-01, -4.8589e-01,  5.5977e-02, -2.9478e-02,\n",
-       "          8.5718e-02,  4.7685e-01,  4.8351e-01, -2.8142e-01],\n",
-       "        [ 1.3377e-01,  5.4434e-01,  3.4505e-01,  1.1307e-01,  4.4057e-01,\n",
-       "         -7.6075e-03,  1.3841e-01, -1.1497e-01, -1.3177e-01,  8.0254e-01,\n",
-       "         -3.0627e-01, -6.8437e-01,  1.9035e-01, -1.0208e+00, -1.3259e-01,\n",
-       "          5.3231e-01, -4.7814e-01, -5.1266e-01,  2.4646e-02, -3.0552e-01,\n",
-       "          2.7398e-01,  5.8269e-01,  6.5481e-01, -4.2041e-01],\n",
-       "        [ 1.9604e-01,  4.0597e-01,  1.9071e-01, -2.5535e-01,  1.1915e-01,\n",
-       "         -6.7129e-02,  5.4386e-03, -8.2196e-02, -4.2803e-01,  7.0287e-01,\n",
-       "         -3.0026e-01, -7.6001e-01, -5.1471e-03, -7.0283e-01, -9.2978e-02,\n",
-       "          1.2243e-01, -1.8398e-01, -4.7374e-01,  2.7978e-01, -3.6962e-01,\n",
-       "          5.6046e-02,  4.1773e-01,  4.9894e-01, -3.1945e-01],\n",
-       "        [ 1.2657e-01,  3.3224e-01,  6.2830e-02,  1.5718e-01,  4.8844e-01,\n",
-       "         -1.1476e-01, -1.5044e-01,  2.5265e-02, -2.0351e-01,  5.5770e-01,\n",
-       "         -3.6036e-01, -7.4406e-01,  1.6962e-01, -9.6185e-01, -2.9334e-01,\n",
-       "          2.2584e-01, -4.1169e-01, -5.2146e-01,  2.3314e-01, -1.3668e-01,\n",
-       "         -1.9598e-02,  3.8727e-01,  3.6892e-01, -3.3071e-01],\n",
-       "        [ 5.2178e-01,  6.9704e-01,  5.0093e-01,  1.1157e-01,  8.0012e-02,\n",
-       "          3.6931e-01, -6.4927e-02,  1.1126e-01, -2.5117e-01,  5.3017e-01,\n",
-       "         -2.6488e-01, -8.4056e-01,  2.2374e-01, -6.6831e-01, -1.9402e-01,\n",
-       "          7.4174e-02, -4.7763e-01, -2.6912e-01,  5.1009e-01, -5.4239e-01,\n",
-       "          3.0123e-01,  3.7529e-01,  4.1625e-01, -2.0141e-01],\n",
-       "        [ 3.7968e-01,  4.9387e-01,  3.6786e-01, -1.3131e-01,  2.4445e-02,\n",
-       "          2.2155e-01, -4.0087e-02, -1.4872e-01, -5.5030e-01,  6.8958e-01,\n",
-       "         -3.8156e-01, -7.5760e-01,  3.2085e-01, -6.4571e-01,  1.1268e-03,\n",
-       "          3.4251e-02, -2.6440e-01, -2.6374e-01,  5.9787e-01, -4.6502e-01,\n",
-       "          2.0074e-01,  4.5471e-01,  2.4238e-01, -4.3247e-01],\n",
-       "        [ 2.9364e-01,  4.8659e-01,  9.0845e-02,  1.6348e-01,  5.7636e-01,\n",
-       "          4.5485e-01, -1.6781e-01, -1.4557e-01, -8.8814e-02,  6.6351e-01,\n",
-       "         -5.3669e-01, -8.2818e-01,  6.0474e-01, -9.4558e-01, -3.0133e-01,\n",
-       "          3.0310e-01, -5.2493e-01, -2.5948e-01,  1.5857e-01, -4.2695e-01,\n",
-       "          2.1311e-01,  4.6502e-01,  8.7946e-02, -5.5815e-01],\n",
-       "        [ 9.2208e-02,  2.9731e-01,  3.3849e-01, -5.1049e-02,  2.7834e-01,\n",
-       "         -1.1120e-01,  1.1835e-01,  1.3665e-01, -2.1291e-01,  3.5107e-01,\n",
-       "         -9.8108e-02, -5.0180e-01,  2.9894e-01, -7.7726e-01, -8.1317e-02,\n",
-       "          3.5704e-01, -3.6759e-01, -2.2148e-01,  1.1019e-01, -1.4452e-02,\n",
-       "          1.5092e-02,  3.3405e-01,  1.2765e-01, -4.0411e-01],\n",
-       "        [ 2.8927e-02,  4.4180e-01,  1.0994e-01,  5.6124e-01,  4.7174e-01,\n",
-       "          1.9914e-01, -9.5047e-02,  3.1277e-02, -1.8656e-01,  5.0631e-01,\n",
-       "         -3.4353e-01, -5.7425e-01,  4.3409e-01, -8.3343e-01, -1.1627e-01,\n",
-       "          3.1852e-02, -4.1274e-01, -2.6756e-01,  4.9652e-01, -2.6137e-01,\n",
-       "          2.8559e-02,  3.0587e-01,  3.6717e-01, -4.4303e-01],\n",
-       "        [-1.0741e-01,  1.3539e-01,  1.5746e-01,  2.1208e-01,  6.3745e-01,\n",
-       "         -2.1864e-01, -1.8820e-01,  2.1184e-01, -3.6832e-02,  3.0890e-01,\n",
-       "         -2.4719e-03, -3.3573e-01,  1.8479e-01, -9.2119e-01, -2.3361e-01,\n",
-       "          8.9827e-02, -5.4372e-01, -4.4935e-01,  3.2967e-01, -9.2807e-02,\n",
-       "          9.9241e-02,  4.1705e-01,  2.4728e-01, -4.8119e-01],\n",
-       "        [ 2.8125e-01,  5.3276e-01,  5.0110e-02,  2.0471e-01,  5.7750e-01,\n",
-       "          4.6670e-02, -2.1400e-01,  6.8794e-03, -6.8737e-02,  4.2138e-01,\n",
-       "         -3.1261e-01, -7.3709e-01,  4.2001e-01, -9.9757e-01, -4.8091e-01,\n",
-       "          2.9960e-01, -6.2133e-01, -4.0566e-01,  3.2191e-01, -1.0219e-02,\n",
-       "          1.2901e-01,  3.9601e-01,  1.6291e-01, -3.3871e-01],\n",
-       "        [ 2.9181e-01,  5.5400e-01,  3.0462e-01,  2.2431e-02,  2.8480e-01,\n",
-       "          4.4624e-01, -2.8859e-01, -1.4629e-01, -4.3573e-02,  2.9742e-01,\n",
-       "         -1.0100e-01, -4.3070e-01,  4.6713e-01, -3.7132e-01, -8.6748e-02,\n",
-       "          2.5666e-01, -3.5361e-01, -2.3917e-02,  3.0071e-01, -3.2420e-01,\n",
-       "          1.3375e-01,  3.4475e-01,  3.0642e-01, -4.3496e-01],\n",
-       "        [-7.7723e-04,  2.3828e-01,  2.3124e-01,  4.1347e-01,  6.8455e-01,\n",
-       "         -9.8319e-03,  1.3403e-01,  1.8460e-02, -1.4025e-01,  5.9780e-01,\n",
-       "         -3.7015e-01, -5.7865e-01,  4.9211e-01, -1.1262e+00, -2.1693e-01,\n",
-       "          3.2002e-01, -2.9313e-01, -3.1941e-01,  9.8446e-02, -6.2767e-02,\n",
-       "         -9.8636e-03,  3.5712e-01,  2.8833e-01, -5.3506e-01]], device='cuda:0',\n",
-       "       grad_fn=<AddmmBackward>)"
-      ]
-     },
-     "execution_count": 25,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "t(x)"
+    "tt = t != 3"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 59,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "PosixPath('hej')"
+       "tensor([[ True,  True,  True, False, False, False]])"
       ]
      },
-     "execution_count": 13,
+     "execution_count": 59,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "t.h"
+    "tt"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 43,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "16"
-      ]
-     },
-     "execution_count": 12,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
-    "t.batch_size"
+    "t = torch.cat((t, t))"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 44,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "PosixPath('hej')"
+       "torch.Size([2, 6])"
       ]
      },
-     "execution_count": 11,
+     "execution_count": 44,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "t.h"
+    "t.shape"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "../text_recognizer/__init__.py\n",
-      "../text_recognizer/callbacks/__init__.py\n",
-      "../text_recognizer/callbacks/wandb_callbacks.py\n",
-      "../text_recognizer/data/image_utils.py\n",
-      "../text_recognizer/data/emnist.py\n",
-      "../text_recognizer/data/iam_lines.py\n",
-      "../text_recognizer/data/download_utils.py\n",
-      "../text_recognizer/data/mappings.py\n",
-      "../text_recognizer/data/iam_preprocessor.py\n",
-      "../text_recognizer/data/__init__.py\n",
-      "../text_recognizer/data/make_wordpieces.py\n",
-      "../text_recognizer/data/iam_paragraphs.py\n",
-      "../text_recognizer/data/sentence_generator.py\n",
-      "../text_recognizer/data/emnist_lines.py\n",
-      "../text_recognizer/data/build_transitions.py\n",
-      "../text_recognizer/data/base_dataset.py\n",
-      "../text_recognizer/data/base_data_module.py\n",
-      "../text_recognizer/data/iam.py\n",
-      "../text_recognizer/data/iam_synthetic_paragraphs.py\n",
-      "../text_recognizer/data/transforms.py\n",
-      "../text_recognizer/data/iam_extended_paragraphs.py\n",
-      "../text_recognizer/networks/__init__.py\n",
-      "../text_recognizer/networks/util.py\n",
-      "../text_recognizer/networks/cnn_tranformer.py\n",
-      "../text_recognizer/networks/encoders/__init__.py\n",
-      "../text_recognizer/networks/encoders/efficientnet/efficientnet.py\n",
-      "../text_recognizer/networks/encoders/efficientnet/__init__.py\n",
-      "../text_recognizer/networks/encoders/efficientnet/utils.py\n",
-      "../text_recognizer/networks/encoders/efficientnet/mbconv.py\n",
-      "../text_recognizer/networks/loss/__init__.py\n",
-      "../text_recognizer/networks/loss/label_smoothing_loss.py\n",
-      "../text_recognizer/networks/vqvae/__init__.py\n",
-      "../text_recognizer/networks/vqvae/decoder.py\n",
-      "../text_recognizer/networks/vqvae/vqvae.py\n",
-      "../text_recognizer/networks/vqvae/vector_quantizer.py\n",
-      "../text_recognizer/networks/vqvae/encoder.py\n",
-      "../text_recognizer/networks/transformer/__init__.py\n",
-      "../text_recognizer/networks/transformer/layers.py\n",
-      "../text_recognizer/networks/transformer/residual.py\n",
-      "../text_recognizer/networks/transformer/attention.py\n",
-      "../text_recognizer/networks/transformer/transformer.py\n",
-      "../text_recognizer/networks/transformer/vit.py\n",
-      "../text_recognizer/networks/transformer/mlp.py\n",
-      "../text_recognizer/networks/transformer/norm.py\n",
-      "../text_recognizer/networks/transformer/positional_encodings/positional_encoding.py\n",
-      "../text_recognizer/networks/transformer/positional_encodings/__init__.py\n",
-      "../text_recognizer/networks/transformer/positional_encodings/absolute_embedding.py\n",
-      "../text_recognizer/networks/transformer/positional_encodings/rotary_embedding.py\n",
-      "../text_recognizer/networks/transformer/nystromer/__init__.py\n",
-      "../text_recognizer/networks/transformer/nystromer/nystromer.py\n",
-      "../text_recognizer/networks/transformer/nystromer/attention.py\n",
-      "../text_recognizer/models/__init__.py\n",
-      "../text_recognizer/models/base.py\n",
-      "../text_recognizer/models/vqvae.py\n",
-      "../text_recognizer/models/transformer.py\n",
-      "../text_recognizer/models/dino.py\n",
-      "../text_recognizer/models/metrics.py\n"
-     ]
-    }
-   ],
-   "source": [
-    "for f in Path(\"../text_recognizer\").glob(\"**/*.py\"):\n",
-    "    print(f)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 45,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "<generator object Path.glob at 0x7ff8bb9ce5f0>"
+       "tensor([[[[ True, False, False, False, False, False],\n",
+       "          [ True,  True, False, False, False, False],\n",
+       "          [ True,  True,  True, False, False, False],\n",
+       "          [ True,  True,  True, False, False, False],\n",
+       "          [ True,  True,  True, False, False, False],\n",
+       "          [ True,  True,  True, False, False, False]]],\n",
+       "\n",
+       "\n",
+       "        [[[ True, False, False, False, False, False],\n",
+       "          [ True,  True, False, False, False, False],\n",
+       "          [ True,  True,  True, False, False, False],\n",
+       "          [ True,  True,  True, False, False, False],\n",
+       "          [ True,  True,  True, False, False, False],\n",
+       "          [ True,  True,  True, False, False, False]]]])"
       ]
      },
-     "execution_count": 12,
+     "execution_count": 45,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "Path(\"..\").glob(\"**/*.py\")"
+    "target_padding_mask(t, 3)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "from text_recognizer.networks.encoders.efficientnet.efficientnet import EfficientNet"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {
-    "scrolled": false
-   },
-   "outputs": [],
-   "source": [
-    "en = EfficientNet(\"b0\")"
+    "target_padding_mask()"
    ]
   },
   {
@@ -1404,7 +1171,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.5"
+   "version": "3.9.6"
   }
  },
  "nbformat": 4,
diff --git a/noxfile.py b/noxfile.py
index a90d53b..d14fefb 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -2,7 +2,6 @@
 import tempfile
 from typing import Any
 
-
 import nox
 from nox.sessions import Session
 
diff --git a/text_recognizer/data/iam_preprocessor.py b/text_recognizer/data/iam_preprocessor.py
index 506036e..f7457e4 100644
--- a/text_recognizer/data/iam_preprocessor.py
+++ b/text_recognizer/data/iam_preprocessor.py
@@ -47,6 +47,8 @@ def load_metadata(
 class Preprocessor:
     """A preprocessor for the IAM dataset."""
 
+    # TODO: attrs
+
     def __init__(
         self,
         data_dir: Union[str, Path],
diff --git a/text_recognizer/networks/cnn_tranformer.py b/text_recognizer/networks/cnn_tranformer.py
index 5c13e9a..e030cb8 100644
--- a/text_recognizer/networks/cnn_tranformer.py
+++ b/text_recognizer/networks/cnn_tranformer.py
@@ -3,6 +3,7 @@ import math
 from typing import Tuple, Type
 
 import attr
+import torch
 from torch import nn, Tensor
 
 from text_recognizer.data.mappings import AbstractMapping
@@ -18,13 +19,19 @@ class CnnTransformer(nn.Module):
     def __attrs_pre_init__(self) -> None:
         super().__init__()
 
-    # Parameters,
+    # Parameters and placeholders,
     input_dims: Tuple[int, int, int] = attr.ib()
     hidden_dim: int = attr.ib()
     dropout_rate: float = attr.ib()
     max_output_len: int = attr.ib()
     num_classes: int = attr.ib()
     padding_idx: int = attr.ib()
+    start_token: str = attr.ib()
+    start_index: int = attr.ib(init=False, default=None)
+    end_token: str = attr.ib()
+    end_index: int = attr.ib(init=False, default=None)
+    pad_token: str = attr.ib()
+    pad_index: int = attr.ib(init=False, default=None)
 
     # Modules.
     encoder: Type[nn.Module] = attr.ib()
@@ -38,6 +45,9 @@ class CnnTransformer(nn.Module):
 
     def __attrs_post_init__(self) -> None:
         """Post init configuration."""
+        self.start_index = int(self.mapping.get_index(self.start_token))
+        self.end_index = int(self.mapping.get_index(self.end_token))
+        self.pad_index = int(self.mapping.get_index(self.pad_token))
         # Latent projector for down sampling number of filters and 2d
         # positional encoding.
         self.latent_encoder = nn.Sequential(
@@ -99,20 +109,20 @@ class CnnTransformer(nn.Module):
         z = self.encoder(x)
         z = self.latent_encoder(z)
 
-        # Permute tensor from [B, E, Ho * Wo] to [Sx, B, E]
-        z = z.permute(2, 0, 1)
+        # Permute tensor from [B, E, Ho * Wo] to [B, Sx, E]
+        z = z.permute(0, 2, 1)
         return z
 
-    def decode(self, z: Tensor, trg: Tensor) -> Tensor:
+    def decode(self, z: Tensor, context: Tensor) -> Tensor:
         """Decodes latent images embedding into word pieces.
 
         Args:
             z (Tensor): Latent images embedding.
-            trg (Tensor): Word embeddings.
+            context (Tensor): Word embeddings.
 
         Shapes:
             - z: :math: `(B, Sx, E)`
-            - trg: :math: `(B, Sy)`
+            - context: :math: `(B, Sy)`
             - out: :math: `(B, Sy, T)`
 
             where Sy is the length of the output and T is the number of tokens.
@@ -120,32 +130,69 @@ class CnnTransformer(nn.Module):
         Returns:
             Tensor: Sequence of word piece embeddings.
         """
-        trg_mask = trg != self.padding_idx
-        trg = self.token_embedding(trg) * math.sqrt(self.hidden_dim)
-        trg = self.token_pos_encoder(trg)
-        out = self.decoder(x=trg, context=z, mask=trg_mask)
+        context_mask = context != self.padding_idx
+        context = self.token_embedding(context) * math.sqrt(self.hidden_dim)
+        context = self.token_pos_encoder(context)
+        out = self.decoder(x=context, context=z, mask=context_mask)
         logits = self.head(out)
         return logits
 
-    def forward(self, x: Tensor, trg: Tensor) -> Tensor:
+    def forward(self, x: Tensor, context: Tensor) -> Tensor:
         """Encodes images into word piece logtis.
 
         Args:
             x (Tensor): Input image(s).
-            trg (Tensor): Target word embeddings.
+            context (Tensor): Target word embeddings.
 
         Shapes:
             - x: :math: `(B, C, H, W)`
-            - trg: :math: `(B, Sy, T)`
+            - context: :math: `(B, Sy, T)`
 
             where B is the batch size, C is the number of input channels, H is
             the image height and W is the image width.
+
+        Returns:
+            Tensor: Sequence of logits.
         """
         z = self.encode(x)
-        logits = self.decode(z, trg)
+        logits = self.decode(z, context)
         return logits
 
     def predict(self, x: Tensor) -> Tensor:
-        """Predicts text in image."""
-        # TODO: continue here!!!!!!!!!
-        pass
+        """Predicts text in image.
+        
+        Args:
+            x (Tensor): Image(s) to extract text from.
+
+        Shapes:
+            - x: :math: `(B, H, W)`
+            - output: :math: `(B, S)`
+
+        Returns:
+            Tensor: A tensor of token indices of the predictions from the model.
+        """
+        bsz = x.shape[0]
+
+        # Encode image(s) to latent vectors.
+        z = self.encode(x)
+
+        # Create a placeholder matrix for storing outputs from the network
+        output = torch.ones((bsz, self.max_output_len), dtype=torch.long).to(x.device)
+        output[:, 0] = self.start_index
+
+        for i in range(1, self.max_output_len):
+            context = output[:, :i]  # (bsz, i)
+            logits = self.decode(z, context)  # (i, bsz, c)
+            tokens = torch.argmax(logits, dim=-1)  # (i, bsz)
+            output[:, i : i + 1] = tokens[-1:]
+
+            # Early stopping of prediction loop if token is end or padding token.
+            if (output[:, i - 1] == self.end_index | output[: i - 1] == self.pad_index).all():
+                break
+
+        # Set all tokens after end token to pad token.
+        for i in range(1, self.max_output_len):
+            idx = (output[:, i -1] == self.end_index | output[:, i - 1] == self.pad_index)
+            output[idx, i] = self.pad_index
+
+        return output
diff --git a/text_recognizer/networks/encoders/efficientnet/efficientnet.py b/text_recognizer/networks/encoders/efficientnet/efficientnet.py
index 59598b5..6719efb 100644
--- a/text_recognizer/networks/encoders/efficientnet/efficientnet.py
+++ b/text_recognizer/networks/encoders/efficientnet/efficientnet.py
@@ -41,7 +41,7 @@ class EfficientNet(nn.Module):
         self.bn_momentum = bn_momentum
         self.bn_eps = bn_eps
         self._conv_stem: nn.Sequential = None
-        self._blocks: nn.Sequential = None
+        self._blocks: nn.ModuleList = None
         self._conv_head: nn.Sequential = None
         self._build()