{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Image Captioning with Transformers\n\nThis notebook is an example on how to generate text from images using transformers. For the encoder part, I draw inspiration from [ViT](https://arxiv.org/abs/2010.11929), while the decoder is an autoregressive generator similar to GPT. To illustrate the model, we will generate captions using the toy dataset MNIST. Instead of classifying the digits, the name will be generated."},{"metadata":{"trusted":true},"cell_type":"code","source":"import pytorch_lightning as pl\nimport torch\nimport matplotlib.pyplot as plt\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom sklearn.datasets import fetch_openml\nimport numpy as np\nfrom torch.utils.data import DataLoader","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The `Dataset` class will give us pairs of images and the corresponding caption. A `1` is added at the beggining of the caption and a `2`is appended at the end. We will use this values for training, indicating the start and end of the caption, respectively. For simplicity, we pad the captions to the maximum length to be able to batch examples."},{"metadata":{"trusted":true},"cell_type":"code","source":"class Dataset(torch.utils.data.Dataset):\n    def __init__(self, X, y, max_len=8):\n        self.X = X \n        self.y = y \n        self.max_len = max_len\n\n    def __len__(self):\n        return len(self.X)\n\n    def __getitem__(self, ix):\n        img = torch.tensor(self.X[ix]).float().view(1, 28, 28)\n        # use 1 for start of sentence\n        # use 2 for end of sentence\n        caption = torch.tensor([1] + self.y[ix] + [2]).long()\n        # use 0 for pad\n        caption = F.pad(caption, (0, self.max_len - len(caption)), 'constant', 0)\n        return img, caption","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The `DataModule` will handle data preparation, tokenization and caption generation. We use a small vocab with all the possible letters in lower case. Each label in the MNIST dataset is replaced by the name of the digit, in Spanish. "},{"metadata":{"trusted":true},"cell_type":"code","source":"class MNISTDataModule(pl.LightningDataModule):\n\n    def __init__(self, batch_size = 1024):\n        super().__init__()\n        self.batch_size = batch_size\n        self.vocab = 'abcdefghijklmnopqrstuvwxyz'\n        self.len_vocab = len(self.vocab) + 3\n\n    def number2caption(self, ix):\n        if ix == 0: return 'cero'\n        if ix == 1: return 'uno'\n        if ix == 2: return 'dos'\n        if ix == 3: return 'tres'\n        if ix == 4: return 'cuatro'\n        if ix == 5: return 'cinco'\n        if ix == 6: return 'seis'\n        if ix == 7: return 'siete'\n        if ix == 8: return 'ocho'\n        if ix == 9: return 'nueve'\n        \n    def caption2ixs(self, caption):\n        return [self.vocab.index(c) + 3 for c in caption]\n\n    def ixs2caption(self, ixs):\n        return ('').join([self.vocab[ix - 3] for ix in ixs if ix not in [0, 1, 2]])\n\n    def setup(self, stage=None):\n        mnist = fetch_openml('mnist_784', version=1)\n        X, y = mnist[\"data\"].values, mnist[\"target\"].values.astype(np.int)\n        # generate captions\n        captions = [self.number2caption(ix) for ix in y]\n        encoded = [self.caption2ixs(caption) for caption in captions]\n        # train / val splits\n        X_train, X_test, y_train, y_test = X[:60000] / 255., X[60000:] / 255., encoded[:60000], encoded[60000:]\n        self.train_ds = Dataset(X_train, y_train)\n        self.val_ds = Dataset(X_test, y_test)\n\n    def train_dataloader(self):\n        return DataLoader(self.train_ds, batch_size=self.batch_size, shuffle=True, pin_memory=True)\n\n    def val_dataloader(self):\n        return DataLoader(self.val_ds, batch_size=self.batch_size, pin_memory=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dm = MNISTDataModule(batch_size=64)\ndm.setup()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"imgs, captions = next(iter(dm.train_dataloader()))\nimgs.shape, captions.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"r, c = 8, 8\nfig = plt.figure(figsize=(c*2, r*2))\nfor _r in range(r):\n    for _c in range(c):\n        ix = _r*c + _c\n        ax = plt.subplot(r, c, ix + 1)\n        img, caption = imgs[ix], captions[ix]\n        ax.axis(\"off\")\n        ax.imshow(img.squeeze(0), cmap=\"gray\")\n        label = dm.ixs2caption(caption)\n        ax.set_title(label)\nplt.tight_layout()\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Now into the model. The inputs to the encoder will be the images re-project, in a patched fashion, to the embedding dimension. You can find more about this in the [ViT](https://arxiv.org/abs/2010.11929) paper."},{"metadata":{"trusted":true},"cell_type":"code","source":"# https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py\n\nclass PatchEmbedding(nn.Module):\n    def __init__(self, img_size, patch_size, in_chans, embed_dim):\n        super().__init__()\n        self.img_size = img_size\n        self.patch_size = patch_size\n        self.n_patches = (img_size // patch_size) ** 2\n        self.patch_size = patch_size\n        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)\n\n    def forward(self, x):\n        x = self.proj(x)  # (B, E, P, P)\n        x = x.flatten(2)  # (B, E, N)\n        x = x.transpose(1, 2)  # (B, N, E)\n        return x","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The model is a Pytorch Transformer that takes images and generates captions. Each image is tiled and reprojected as mentioned before, and fed to the encoder. The decoder is masked so it can only attend to past tokens. Finally, we use a linear classifier on top of the decoder to output a probability distribution over the vocabulary for each generated word."},{"metadata":{"trusted":true},"cell_type":"code","source":"class Model(pl.LightningModule):\n\n    def __init__(self, \n                 len_vocab,\n                 img_size=28, \n                 patch_size=7, \n                 in_chans=1, \n                 embed_dim=100, \n                 max_len=8, \n                 nhead=2, \n                 num_encoder_layers=3,\n                 num_decoder_layers=3,\n                 dim_feedforward=400,\n                 dropout=0.1\n                ):\n        super().__init__()\n        \n        self.patch_embed = PatchEmbedding(img_size, patch_size, in_chans, embed_dim)\n        self.pos_embed = nn.Parameter(torch.zeros(1, self.patch_embed.n_patches, embed_dim))\n        \n        self.trg_emb = nn.Embedding(len_vocab, embed_dim)\n        self.trg_pos_emb = nn.Embedding(max_len, embed_dim)\n        self.max_len = max_len\n\n        self.transformer = torch.nn.Transformer(\n            embed_dim, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, dropout\n        )\n        \n        self.l = nn.LayerNorm(embed_dim)\n        self.fc = nn.Linear(embed_dim, len_vocab)\n\n    def forward(self, images, captions):\n        # embed images\n        embed_imgs = self.patch_embed(images)\n        embed_imgs = embed_imgs + self.pos_embed  \n        # embed captions\n        B, trg_seq_len = captions.shape \n        trg_positions = (torch.arange(0, trg_seq_len).expand(B, trg_seq_len).to(self.device))\n        embed_trg = self.trg_emb(captions) + self.trg_pos_emb(trg_positions)\n        trg_mask = self.transformer.generate_square_subsequent_mask(trg_seq_len).to(self.device)\n        tgt_padding_mask = captions == 0\n        # transformer\n        y = self.transformer(\n            embed_imgs.permute(1,0,2),  \n            embed_trg.permute(1,0,2),  \n            tgt_mask=trg_mask, \n            tgt_key_padding_mask = tgt_padding_mask\n        ).permute(1,0,2) \n        # head\n        return self.fc(self.l(y))\n\n    def predict(self, images):\n        self.eval()\n        with torch.no_grad():\n            images = images.to(self.device)\n            B = images.shape[0]\n            eos = torch.tensor([1], dtype=torch.long, device=self.device).expand(B, 1)\n            trg_input = eos\n            for _ in range(self.max_len):\n                preds = self(images, trg_input)\n                preds = torch.argmax(preds, axis=2)\n                trg_input = torch.cat([eos, preds], 1)\n            return preds\n        \n    def compute_loss_and_acc(self, batch):\n        x, y = batch\n        y_hat = self(x, y[:,:-1])\n        trg_output = y[:,1:] \n        loss = F.cross_entropy(y_hat.permute(0,2,1), trg_output) \n        # I know this is not the best metric...\n        acc = (torch.argmax(y_hat, axis=2) == trg_output).sum().item() / (trg_output.shape[0]*trg_output.shape[1])\n        return loss, acc\n    \n    def training_step(self, batch, batch_idx):\n        loss, acc = self.compute_loss_and_acc(batch)\n        self.log('loss', loss)\n        self.log('acc', acc, prog_bar=True)\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        loss, acc = self.compute_loss_and_acc(batch)\n        self.log('val_loss', loss, prog_bar=True)\n        self.log('val_acc', acc, prog_bar=True)\n\n    def configure_optimizers(self):\n        optimizer = torch.optim.Adam(self.parameters(), lr=0.001)\n        return optimizer","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We now train the model"},{"metadata":{"trusted":true},"cell_type":"code","source":"dm = MNISTDataModule()\nmodel = Model(dm.len_vocab)\ntrainer = pl.Trainer(max_epochs=10, gpus=1, precision=16)\ntrainer.fit(model, dm)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"and generate some captions for the validation images."},{"metadata":{"trusted":true},"cell_type":"code","source":"imgs, captions = next(iter(dm.val_dataloader()))\npreds = model.predict(imgs)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"r, c = 8, 8\nfig = plt.figure(figsize=(c*2, r*2))\nfor _r in range(r):\n    for _c in range(c):\n        ix = _r*c + _c\n        ax = plt.subplot(r, c, ix + 1)\n        img, caption = imgs[ix], captions[ix]\n        ax.axis(\"off\")\n        ax.imshow(img.squeeze(0), cmap=\"gray\")\n        label = dm.ixs2caption(caption)\n        pred = dm.ixs2caption(preds[ix])\n        ax.set_title(label)\n        ax.set_title(f'{label}/{pred}', color=\"green\" if label == pred else 'red')\nplt.tight_layout()\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"I hope this notebook is useful for you. I wanted to use this challenge as an excuse to learn about transformers. If you find a bug or something that can be improved, please let me know ! I think this notebook can easily be extended to work with molecular imagery and generate InChIs :)"}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}