24. PyTorch's-Transformer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# utils\n",
    "from utils import count_parameters\n",
    "import torch\n",
    "import math\n",
    "\n",
    "# data\n",
    "from torchtext.datasets import Multi30k\n",
    "from torchtext.data import Field, BucketIterator\n",
    "\n",
    "# model\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "# training\n",
    "import torch.optim as optim\n",
    "import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cpu')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create data fields for source and target\n",
    "source = Field(\n",
    "    init_token=\"<sos>\",\n",
    "    eos_token=\"<eos>\",\n",
    "    lower=True,\n",
    "    tokenize=\"spacy\",\n",
    "    tokenizer_language=\"de\",\n",
    "    batch_first=True\n",
    ")\n",
    "target = Field(\n",
    "    init_token=\"<sos>\",\n",
    "    eos_token=\"<eos>\",\n",
    "    lower=True,\n",
    "    tokenize=\"spacy\",\n",
    "    tokenizer_language=\"de\",\n",
    "    batch_first=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# download the parallel corpus\n",
    "train, val, test = Multi30k.splits(\n",
    "    exts=(\".de\", \".en\"),\n",
    "    fields=(source, target)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build the vocab\n",
    "source.build_vocab(train)\n",
    "target.build_vocab(train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create data loaders\n",
    "BATCH_SIZE = 64\n",
    "train_loader, val_loader, test_loader = BucketIterator.splits(\n",
    "    datasets=(train, val, test),\n",
    "    batch_size=BATCH_SIZE,\n",
    "    device=device,\n",
    "    shuffle=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([64, 23]) torch.Size([64, 26])\n"
     ]
    }
   ],
   "source": [
    "batch =  next(iter(train_loader))\n",
    "print(batch.src.shape, batch.trg.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PyTorch's Transformer's Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EmbeddingLayer(nn.Module):\n",
    "    \"\"\"\n",
    "        takes input as token and convert it into embeddings\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, vocab_size, embedding_dim):\n",
    "        \n",
    "        super(EmbeddingLayer, self).__init__()  \n",
    "        \n",
    "        self.vocab_size = vocab_size\n",
    "        self.embedding_dim = embedding_dim\n",
    "        \n",
    "        self.embedding = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embedding_dim)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "            Ideally x.shape -> [batch, seq_len]\n",
    "        \"\"\"\n",
    "        embedded = self.embedding(x)\n",
    "        return embedded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "# encoder model\n",
    "class Encoder(nn.Module):\n",
    "    \"\"\"\n",
    "        encoder will get source token and will produce contextualized embedding\n",
    "    \"\"\"\n",
    "    def __init__(self, embedding_dim, vocab_size, num_heads=8, num_layers=1, max_len=100, dropout=0.15):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.embedding_dim = embedding_dim\n",
    "        self.vocab_size = vocab_size\n",
    "        \n",
    "        # same embedding can be used for both sys\n",
    "        self.tok_embedding = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embedding_dim)\n",
    "        self.pos_embedding = nn.Embedding(num_embeddings=max_len, embedding_dim=embedding_dim)\n",
    "        \n",
    "        # scaling\n",
    "        self.scale = torch.sqrt(torch.FloatTensor([embedding_dim])).to(device)\n",
    "    \n",
    "        self.encoder = nn.TransformerEncoder(\n",
    "            encoder_layer=nn.TransformerEncoderLayer(d_model=embedding_dim, nhead=num_heads), \n",
    "            num_layers=num_layers, \n",
    "            norm=nn.LayerNorm(normalized_shape=embedding_dim)\n",
    "        )\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "    \n",
    "    \n",
    "    def forward(self, src, src_mask=None):\n",
    "        \n",
    "        \"\"\"\n",
    "            src.shape -> [batch, src_len]\n",
    "        \"\"\"\n",
    "        batch, src_len = src.shape[0], src.shape[1]\n",
    "        \n",
    "        # create position tensor, shape will be [batch, src_len] by dooing so batch_first will be True\n",
    "        position  = torch.arange(start=0, end=src_len, device=device).unsqueeze(0).repeat(batch, 1)\n",
    "        \n",
    "        # embeddings\n",
    "        tok_embedded = self.tok_embedding(src)\n",
    "        pos_embedded = self.pos_embedding(position)\n",
    "        \n",
    "        # scale the token embeddings by multiplyig it with srqt(d_model) where d_model is embedding_dim\n",
    "        tok_scaled = tok_embedded * self.scale\n",
    "        \n",
    "        # add the scaled_tok and position embedding and then apply dropout, that will be input to the encoder\n",
    "        encoder_input = self.dropout(tok_scaled + pos_embedded)\n",
    "        \n",
    "        \n",
    "        encoded = self.encoder(encoder_input)\n",
    "        \n",
    "        return encoded\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    \n",
    "    def __init__(self, embedding_dim, vocab_size, num_heads=8, num_layers=8, max_len=100, dropout=0.2):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.embedding_dim = embedding_dim\n",
    "        self.vocab_size = vocab_size\n",
    "        \n",
    "        self.tok_embedding = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embedding_dim)\n",
    "        self.pos_embedding = nn.Embedding(num_embeddings=max_len, embedding_dim=embedding_dim)\n",
    "        \n",
    "        # scaling\n",
    "        self.scale = torch.sqrt(torch.FloatTensor([embedding_dim])).to(device)\n",
    "        \n",
    "        # dropout layer\n",
    "        self.dropout = nn.Dropout(p=0.15)\n",
    "        \n",
    "        self.decoder = nn.TransformerDecoder(\n",
    "            decoder_layer=nn.TransformerDecoderLayer(d_model=embedding_dim, nhead=num_heads),\n",
    "            num_layers=num_layers,\n",
    "            norm=nn.LayerNorm(normalized_shape=embedding_dim)\n",
    "        )\n",
    "        \n",
    "        \n",
    "        self.fc_out = nn.Linear(in_features=embedding_dim, out_features=vocab_size)\n",
    "    \n",
    "    def forward(self, trg, src_encoded):\n",
    "        \n",
    "        \"\"\"\n",
    "            trg.shape -> [batch, trg_len]\n",
    "        \"\"\"\n",
    "        batch, trg_len = trg.shape[0], trg.shape[1]\n",
    "        \n",
    "        # create position tensor, shape will be [batch, src_len] by dooing so batch_first will be True\n",
    "        position  = torch.arange(start=0, end=trg_len, device=device).unsqueeze(0).repeat(batch, 1)\n",
    "        \n",
    "        # embeddings\n",
    "        tok_embedded = self.tok_embedding(trg)\n",
    "        pos_embedded = self.pos_embedding(position)\n",
    "        \n",
    "        # scale the token embeddings by multiplyig it with srqt(d_model) where d_model is embedding_dim\n",
    "        tok_scaled = tok_embedded * self.scale\n",
    "        \n",
    "        # add the scaled_tok and position embedding and then apply dropout, that will be input to the encoder\n",
    "        decoder_input = self.dropout(tok_scaled + pos_embedded)\n",
    "        \n",
    "        outputs = self.decoder(decoder_input.permute(1, 0, 2), src_encoded.permute(1, 0, 2))\n",
    "        prediction = self.fc_out(outputs)\n",
    "        return prediction\n",
    "         "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training with PyTorch-Lightning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytorch_lightning as pl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(pl.LightningModule):\n",
    "    \n",
    "    def __init__(self, encoder, decoder, PAD_IDX):\n",
    "        super(Transformer, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "    \n",
    "    def forward(self, src, trg):\n",
    "        src_encoded = self.encoder(src)\n",
    "        outputs = self.decoder(trg, src_encoded)\n",
    "        return outputs\n",
    "    \n",
    "    def configure_optimizers(self):\n",
    "        return optim.Adam(params=self.parameters(), lr=1e-3)\n",
    "    \n",
    "    def train_dataloader(self):\n",
    "        return train_loader\n",
    "    \n",
    "    def training_step(self, batch, batch_idx):\n",
    "        src, trg = batch.src, batch.trg\n",
    "        batch_size, trg_len = trg.shape[0], trg.shape[1] # batch_first was true in BucketIterato\n",
    "        outputs = self(src, trg)\n",
    "        loss = F.cross_entropy(outputs.view(batch_size*trg_len, -1), trg.view(-1), ignore_index=PAD_IDX)\n",
    "        ppl = torch.exp(loss)\n",
    "        tensorboard_logs = {\"loss\":loss, \"ppl\":ppl}\n",
    "        return {\"loss\":loss, \"ppl\":ppl,\"log\":tensorboard_logs}\n",
    "    \n",
    "    def val_dataloader(self):\n",
    "        return val_loader\n",
    "    \n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        src, trg = batch.src, batch.trg\n",
    "        batch_size, trg_len = trg.shape[0], trg.shape[1] # batch_first was true in BucketIterato\n",
    "        outputs = self(src, trg)\n",
    "        loss = F.cross_entropy(outputs.view(batch_size*trg_len, -1), trg.view(-1), ignore_index=PAD_IDX)\n",
    "        ppl = torch.exp(loss)\n",
    "        tensorboard_logs = {\"val_loss\":loss, \"val_ppl\":ppl}\n",
    "        return {\"val_loss\":loss, \"val_ppl\":ppl,\"log\":tensorboard_logs}\n",
    "    \n",
    "    def validation_epoch_end(self, outputs):\n",
    "        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()\n",
    "        avg_ppl = torch.stack([x['val_ppl'] for x in outputs]).mean()\n",
    "        tensorboard_logs = {'val_loss': avg_loss, 'val_ppl':avg_ppl}\n",
    "        return {'val_loss': avg_loss,'val_ppl':avg_ppl, 'log': tensorboard_logs}    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [],
   "source": [
    "src_vocab_size = len(source.vocab)\n",
    "trg_vocab_size = len(target.vocab)\n",
    "PAD_IDX = target.vocab.stoi[target.pad_token]\n",
    "embedding_dim = 256"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder = Encoder(embedding_dim=embedding_dim, vocab_size=src_vocab_size)\n",
    "decoder = Decoder(embedding_dim=embedding_dim, vocab_size=trg_vocab_size)\n",
    "transformer = Transformer(encoder=encoder, decoder=decoder, PAD_IDX=PAD_IDX)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: False, used: False\n",
      "TPU available: False, using: 0 TPU cores\n"
     ]
    }
   ],
   "source": [
    "trainer = pl.Trainer(max_epochs=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "  | Name    | Type    | Params\n",
      "------------------------------------\n",
      "0 | encoder | Encoder | 6 M   \n",
      "1 | decoder | Decoder | 17 M  \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1:  97%|█████████▋| 454/470 [09:44<00:20,  1.29s/it, loss=5.413, v_num=5]\n",
      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
      "Epoch 1:  97%|█████████▋| 455/470 [09:44<00:19,  1.28s/it, loss=5.413, v_num=5]\n",
      "Epoch 1:  97%|█████████▋| 456/470 [09:44<00:17,  1.28s/it, loss=5.413, v_num=5]\n",
      "Epoch 1:  97%|█████████▋| 457/470 [09:44<00:16,  1.28s/it, loss=5.413, v_num=5]\n",
      "Epoch 1:  97%|█████████▋| 458/470 [09:44<00:15,  1.28s/it, loss=5.413, v_num=5]\n",
      "Epoch 1:  98%|█████████▊| 459/470 [09:44<00:14,  1.27s/it, loss=5.413, v_num=5]\n",
      "Epoch 1:  98%|█████████▊| 460/470 [09:45<00:12,  1.27s/it, loss=5.413, v_num=5]\n",
      "Epoch 1:  98%|█████████▊| 461/470 [09:45<00:11,  1.27s/it, loss=5.413, v_num=5]\n",
      "Epoch 1:  98%|█████████▊| 462/470 [09:45<00:10,  1.27s/it, loss=5.413, v_num=5]\n",
      "Epoch 1:  99%|█████████▊| 463/470 [09:45<00:08,  1.26s/it, loss=5.413, v_num=5]\n",
      "Epoch 1:  99%|█████████▊| 464/470 [09:45<00:07,  1.26s/it, loss=5.413, v_num=5]\n",
      "Epoch 1:  99%|█████████▉| 465/470 [09:46<00:06,  1.26s/it, loss=5.413, v_num=5]\n",
      "Epoch 1:  99%|█████████▉| 466/470 [09:46<00:05,  1.26s/it, loss=5.413, v_num=5]\n",
      "Epoch 1:  99%|█████████▉| 467/470 [09:46<00:03,  1.26s/it, loss=5.413, v_num=5]\n",
      "Epoch 1: 100%|█████████▉| 468/470 [09:46<00:02,  1.25s/it, loss=5.413, v_num=5]\n",
      "Epoch 1: 100%|█████████▉| 469/470 [09:47<00:01,  1.25s/it, loss=5.413, v_num=5]\n",
      "Epoch 1: 100%|██████████| 470/470 [09:47<00:00,  1.25s/it, loss=5.413, v_num=5]\n",
      "Epoch 2:  97%|█████████▋| 454/470 [09:29<00:20,  1.25s/it, loss=5.368, v_num=5]\n",
      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
      "Epoch 2:  97%|█████████▋| 455/470 [09:29<00:18,  1.25s/it, loss=5.368, v_num=5]\n",
      "Epoch 2:  97%|█████████▋| 456/470 [09:29<00:17,  1.25s/it, loss=5.368, v_num=5]\n",
      "Epoch 2:  97%|█████████▋| 457/470 [09:30<00:16,  1.25s/it, loss=5.368, v_num=5]\n",
      "Epoch 2:  97%|█████████▋| 458/470 [09:30<00:14,  1.25s/it, loss=5.368, v_num=5]\n",
      "Epoch 2:  98%|█████████▊| 459/470 [09:30<00:13,  1.24s/it, loss=5.368, v_num=5]\n",
      "Epoch 2:  98%|█████████▊| 460/470 [09:30<00:12,  1.24s/it, loss=5.368, v_num=5]\n",
      "Epoch 2:  98%|█████████▊| 461/470 [09:30<00:11,  1.24s/it, loss=5.368, v_num=5]\n",
      "Epoch 2:  98%|█████████▊| 462/470 [09:31<00:09,  1.24s/it, loss=5.368, v_num=5]\n",
      "Epoch 2:  99%|█████████▊| 463/470 [09:31<00:08,  1.23s/it, loss=5.368, v_num=5]\n",
      "Epoch 2:  99%|█████████▊| 464/470 [09:31<00:07,  1.23s/it, loss=5.368, v_num=5]\n",
      "Epoch 2:  99%|█████████▉| 465/470 [09:31<00:06,  1.23s/it, loss=5.368, v_num=5]\n",
      "Epoch 2:  99%|█████████▉| 466/470 [09:32<00:04,  1.23s/it, loss=5.368, v_num=5]\n",
      "Epoch 2:  99%|█████████▉| 467/470 [09:32<00:03,  1.23s/it, loss=5.368, v_num=5]\n",
      "Epoch 2: 100%|█████████▉| 468/470 [09:32<00:02,  1.22s/it, loss=5.368, v_num=5]\n",
      "Epoch 2: 100%|█████████▉| 469/470 [09:32<00:01,  1.22s/it, loss=5.368, v_num=5]\n",
      "Epoch 2: 100%|██████████| 470/470 [09:33<00:00,  1.22s/it, loss=5.368, v_num=5]\n",
      "Epoch 2: 100%|██████████| 470/470 [09:33<00:00,  1.22s/it, loss=5.368, v_num=5]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.fit(transformer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
      "Validating:   6%|▋         | 1/16 [00:00<00:03,  4.80it/s]\u001b[A\n",
      "Validating:  12%|█▎        | 2/16 [00:00<00:02,  5.14it/s]\u001b[A\n",
      "Validating:  19%|█▉        | 3/16 [00:00<00:02,  5.48it/s]\u001b[A\n",
      "Validating:  25%|██▌       | 4/16 [00:00<00:02,  5.73it/s]\u001b[A\n",
      "Validating:  31%|███▏      | 5/16 [00:00<00:01,  5.53it/s]\u001b[A\n",
      "Validating:  38%|███▊      | 6/16 [00:01<00:01,  5.15it/s]\u001b[A\n",
      "Validating:  44%|████▍     | 7/16 [00:01<00:01,  5.14it/s]\u001b[A\n",
      "Validating:  50%|█████     | 8/16 [00:01<00:01,  5.08it/s]\u001b[A\n",
      "Validating:  56%|█████▋    | 9/16 [00:01<00:01,  5.37it/s]\u001b[A\n",
      "Validating:  62%|██████▎   | 10/16 [00:01<00:01,  5.58it/s]\u001b[A\n",
      "Validating:  69%|██████▉   | 11/16 [00:02<00:00,  5.60it/s]\u001b[A\n",
      "Validating:  75%|███████▌  | 12/16 [00:02<00:00,  5.43it/s]\u001b[A\n",
      "Validating:  81%|████████▏ | 13/16 [00:02<00:00,  4.93it/s]\u001b[A\n",
      "Validating:  88%|████████▊ | 14/16 [00:02<00:00,  4.90it/s]\u001b[A\n",
      "Validating:  94%|█████████▍| 15/16 [00:02<00:00,  4.69it/s]\u001b[A\n",
      "Validating: 100%|██████████| 16/16 [00:03<00:00,  4.40it/s]\u001b[A\n",
      "                                                           \u001b[A"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'val_loss': tensor(5.3635), 'val_ppl': tensor(216.9025)}"
      ]
     },
     "execution_count": 82,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.run_evaluation()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.6 64-bit ('base': conda)",
   "language": "python",
   "name": "python37664bitbaseconda1b4d65181bfe435290e55078ed6e0090"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}