transformer/transformer_torch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import time\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.functional import F"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Train Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "char_to_id = {\n",
    "    '<PAD>': 0,\n",
    "    '<BOS>': 1,\n",
    "    '<EOS>': 2,\n",
    "    '0': 3,\n",
    "    '1': 4,\n",
    "    '2': 5,\n",
    "    '3': 6,\n",
    "    '4': 7,\n",
    "    '5': 8,\n",
    "    '6': 9,\n",
    "    '7': 10,\n",
    "    '8': 11,\n",
    "    '9': 12,\n",
    "    '+': 13,\n",
    "    '-': 14,\n",
    "}\n",
    "\n",
    "id_to_char = {\n",
    "    0: '<PAD>',\n",
    "    1: '<BOS>',\n",
    "    2: '<EOS>',\n",
    "    3: '0',\n",
    "    4: '1',\n",
    "    5: '2',\n",
    "    6: '3',\n",
    "    7: '4',\n",
    "    8: '5',\n",
    "    9: '6',\n",
    "    10: '7',\n",
    "    11: '8',\n",
    "    12: '9',\n",
    "    13: '+',\n",
    "    14: '-',\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def id_list_to_sequence(sequence):\n",
    "    return ''.join([id_to_char.get(i) for i in sequence])\n",
    "\n",
    "def sequence_to_id_list(sequence):\n",
    "    return [char_to_id.get(c) for c in sequence]\n",
    "\n",
    "def create_dataset(size, num_digit=5, ops=['+', '-']):\n",
    "    source_sequences = []\n",
    "    target_sequences = []\n",
    "\n",
    "    for _ in range(size):\n",
    "        a = random.randint(0, 10**num_digit)\n",
    "        b = random.randint(0, 10**num_digit)\n",
    "        op = random.choice(ops)\n",
    "\n",
    "        if op == '+':\n",
    "            source_tokens = '{}+{}'.format(a, b)\n",
    "            target_tokens = '{}'.format(a + b)\n",
    "        elif op == '-':\n",
    "            source_tokens = '{}-{}'.format(a, b)\n",
    "            target_tokens = '{}'.format(a - b)\n",
    "\n",
    "        source_sequences.append(source_tokens)\n",
    "        target_sequences.append(target_tokens)\n",
    "\n",
    "    return source_sequences, target_sequences\n",
    "\n",
    "\n",
    "def tokenize(sequences, bos=False, eos=False):\n",
    "    tensor = [torch.LongTensor(\n",
    "        ([char_to_id['<BOS>']] if bos else []) + \\\n",
    "        sequence_to_id_list(s) + \\\n",
    "        ([char_to_id['<EOS>']] if eos else []))\n",
    "        for s in sequences\n",
    "    ]\n",
    "    tensor = torch.nn.utils.rnn.pad_sequence(\n",
    "        tensor, batch_first=True, padding_value=0)\n",
    "    return tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data_num = 25600\n",
    "batch_size = 256\n",
    "\n",
    "train_source_sequences, train_target_sequences = create_dataset(train_data_num)\n",
    "data_loader = torch.utils.data.DataLoader(\n",
    "    dataset=torch.utils.data.TensorDataset(\n",
    "        tokenize(train_source_sequences),\n",
    "        tokenize(train_target_sequences, bos=True),\n",
    "        tokenize(train_target_sequences, eos=True)\n",
    "    ), batch_size=batch_size, shuffle=True, drop_last=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(10):\n",
    "    print(train_source_sequences[i], '=', train_target_sequences[i])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Implementation of Transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab_size = len(char_to_id)\n",
    "num_blocks = 2\n",
    "num_hidden_size = 128\n",
    "num_heads = 8\n",
    "dropout_rate = 0.1\n",
    "\n",
    "num_epochs = 100\n",
    "num_batches = train_data_num // batch_size\n",
    "\n",
    "device = torch.cuda.current_device() if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "NEG_INF = -1e9\n",
    "\n",
    "\n",
    "def get_position_encoding(length, hidden_size, device, dtype=torch.float32):\n",
    "    position = torch.arange(length, dtype=dtype, device=device)\n",
    "    timescale = torch.arange(hidden_size // 2, dtype=dtype, device=device)\n",
    "\n",
    "    angle_rates = 1 / (10000 ** ((2 * timescale) / hidden_size))\n",
    "    angle_rads = position[:, None] * angle_rates[None, :]\n",
    "\n",
    "    position_encoding = torch.stack([torch.sin(angle_rads), torch.cos(angle_rads)], axis=2)\n",
    "    position_encoding = position_encoding.view(length, hidden_size)\n",
    "\n",
    "    return position_encoding\n",
    "\n",
    "\n",
    "def get_padding_bias(x):\n",
    "    attention_bias = (x == 0) * NEG_INF\n",
    "    attention_bias = attention_bias[:, None, None, :]\n",
    "    return attention_bias\n",
    "\n",
    "\n",
    "def get_decoder_self_attention_bias(length, device):\n",
    "    valid_locs = torch.tril(torch.ones(length, length, device=device))\n",
    "    valid_locs = valid_locs[None, None, :, :]\n",
    "    decoder_bias = NEG_INF * (1.0 - valid_locs)\n",
    "    return decoder_bias"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EmbeddingSharedWeights(nn.Module):\n",
    "\n",
    "    def __init__(self, vocab_size, hidden_size):\n",
    "        super(EmbeddingSharedWeights, self).__init__()\n",
    "        self.vocab_size = vocab_size\n",
    "        self.hidden_size = hidden_size\n",
    "        self.embedding_weights = nn.Parameter(\n",
    "            torch.normal(mean=0., std=num_hidden_size**-0.5, size=(vocab_size, num_hidden_size))\n",
    "        )\n",
    "\n",
    "    def forward(self, inputs, mode='embedding'):\n",
    "        if mode == 'embedding':\n",
    "            return self._embedding(inputs)\n",
    "        elif mode == 'linear':\n",
    "            return self._linear(inputs)\n",
    "\n",
    "    def _embedding(self, inputs):\n",
    "        embeddings = F.embedding(inputs, self.embedding_weights, padding_idx=0)\n",
    "        embeddings *= self.hidden_size ** 0.5\n",
    "        return embeddings\n",
    "\n",
    "    def _linear(self, inputs):\n",
    "        outputs = inputs @ self.embedding_weights.transpose(1, 0)\n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForwardNetwork(nn.Module):\n",
    "\n",
    "    def __init__(self, hidden_size, filter_size, dropout_rate) -> None:\n",
    "        super(FeedForwardNetwork, self).__init__()\n",
    "        self.filter_dense_layer = nn.Linear(hidden_size, filter_size)\n",
    "        self.output_dense_layer = nn.Linear(filter_size, hidden_size)\n",
    "        self.gelu = nn.GELU()\n",
    "        self.dropout_layer = nn.Dropout(dropout_rate)\n",
    "\n",
    "    def forward(self, x):\n",
    "        output = self.filter_dense_layer(x)\n",
    "        output = self.gelu(output)\n",
    "        output = self.dropout_layer(output)\n",
    "        output = self.output_dense_layer(output)\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Attention(nn.Module):\n",
    "\n",
    "    def __init__(self, hidden_size, num_heads, dropout_rate):\n",
    "        super(Attention, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "        self.num_heads = num_heads\n",
    "        self.depth = hidden_size // num_heads\n",
    "\n",
    "        self.q_dense_layer = nn.Linear(hidden_size, hidden_size, bias=False)\n",
    "        self.k_dense_layer = nn.Linear(hidden_size, hidden_size, bias=False)\n",
    "        self.v_dense_layer = nn.Linear(hidden_size, hidden_size, bias=False)\n",
    "        self.output_dense_layer = nn.Linear(hidden_size, hidden_size, bias=False)\n",
    "        self.dropout_layer = nn.Dropout(dropout_rate)\n",
    "\n",
    "    def forward(self, x, y, bias):\n",
    "        q = self.q_dense_layer(x)\n",
    "        k = self.k_dense_layer(y)\n",
    "        v = self.v_dense_layer(y)\n",
    "\n",
    "        q = self._split_heads(q)\n",
    "        k = self._split_heads(k)\n",
    "        v = self._split_heads(v)\n",
    "\n",
    "        q = q * self.depth ** -0.5\n",
    "\n",
    "        logits = q @ k.transpose(-2, -1)\n",
    "        logits += bias\n",
    "        weights = F.softmax(logits, dim=-1)\n",
    "        weights = self.dropout_layer(weights)\n",
    "\n",
    "        attention_output = weights @ v \n",
    "        attention_output = self._combine_heads(attention_output)\n",
    "        attention_output = self.output_dense_layer(attention_output)\n",
    "\n",
    "        return attention_output\n",
    "\n",
    "    def _split_heads(self, x):\n",
    "        batch_size, length, _ = x.size()\n",
    "        return x.view(batch_size, length, self.num_heads, self.depth).transpose(1, 2)\n",
    "\n",
    "    def _combine_heads(self, x):\n",
    "        batch_size, _, length, _ = x.size()\n",
    "        return x.transpose(1, 2).contiguous().view(batch_size, length, self.hidden_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SelfAttention(Attention):\n",
    "    def __call__(self, x, bias):\n",
    "        return super(SelfAttention, self).__call__(x, x, bias)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PrePostProcessingWrapper(nn.Module):\n",
    "\n",
    "    def __init__(self, layer, hidden_size, dropout_rate):\n",
    "        super(PrePostProcessingWrapper, self).__init__()\n",
    "        self.layer = layer\n",
    "        self.layer_norm = nn.LayerNorm(hidden_size, eps=1e-6)\n",
    "        self.dropout_layer = nn.Dropout(dropout_rate)\n",
    "\n",
    "    def forward(self, x, *args, **kwargs):\n",
    "        y = self.layer_norm(x)\n",
    "        y = self.layer(y, *args, **kwargs)\n",
    "        y = self.dropout_layer(y)\n",
    "\n",
    "        return x + y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderStack(nn.Module):\n",
    "\n",
    "    def __init__(self, num_blocks, hidden_size, num_heads, dropout_rate):\n",
    "        super(EncoderStack, self).__init__()\n",
    "        self.layers = nn.ModuleList()\n",
    "        for _ in range(num_blocks):\n",
    "            self_attention_layer = SelfAttention(hidden_size, num_heads, dropout_rate)\n",
    "            feed_forward_network = FeedForwardNetwork(hidden_size, hidden_size * 4, dropout_rate)\n",
    "\n",
    "            self.layers.append(nn.ModuleList([\n",
    "                PrePostProcessingWrapper(self_attention_layer, hidden_size, dropout_rate),\n",
    "                PrePostProcessingWrapper(feed_forward_network, hidden_size, dropout_rate)\n",
    "            ]))\n",
    "\n",
    "        self.output_normalization = nn.LayerNorm(hidden_size, eps=1e-6)\n",
    "\n",
    "    def forward(self, encoder_inputs, attention_bias):\n",
    "        for n, layer in enumerate(self.layers):\n",
    "            self_attention_layer = layer[0]\n",
    "            feed_forward_network = layer[1]\n",
    "\n",
    "            encoder_inputs = self_attention_layer(encoder_inputs, attention_bias)\n",
    "            encoder_inputs = feed_forward_network(encoder_inputs)\n",
    "\n",
    "        return self.output_normalization(encoder_inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderStack(nn.Module):\n",
    "\n",
    "    def __init__(self, num_blocks, hidden_size, num_heads, dropout_rate):\n",
    "        super(DecoderStack, self).__init__()\n",
    "        self.layers = nn.ModuleList()\n",
    "        for _ in range(num_blocks):\n",
    "            self_attention_layer = SelfAttention(hidden_size, num_heads, dropout_rate)\n",
    "            enc_dec_attention_layer = Attention(hidden_size, num_heads, dropout_rate)\n",
    "            feed_forward_network = FeedForwardNetwork(hidden_size, hidden_size * 4, dropout_rate)\n",
    "\n",
    "            self.layers.append(nn.ModuleList([\n",
    "                PrePostProcessingWrapper(self_attention_layer, hidden_size, dropout_rate),\n",
    "                PrePostProcessingWrapper(enc_dec_attention_layer, hidden_size, dropout_rate),\n",
    "                PrePostProcessingWrapper(feed_forward_network, hidden_size, dropout_rate)\n",
    "            ]))\n",
    "\n",
    "        self.output_normalization = nn.LayerNorm(hidden_size, eps=1e-6)\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        decoder_inputs,\n",
    "        encoder_outputs,\n",
    "        decoder_self_attention_bias,\n",
    "        attention_bias\n",
    "    ):\n",
    "        for n, layer in enumerate(self.layers):\n",
    "            self_attention_layer = layer[0]\n",
    "            enc_dec_attention_layer = layer[1]\n",
    "            feed_forward_network = layer[2]\n",
    "\n",
    "            decoder_inputs = self_attention_layer(decoder_inputs, decoder_self_attention_bias)\n",
    "            decoder_inputs = enc_dec_attention_layer(decoder_inputs, encoder_outputs, attention_bias)\n",
    "            decoder_inputs = feed_forward_network(decoder_inputs)\n",
    "\n",
    "        return self.output_normalization(decoder_inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        vocab_size,\n",
    "        num_blocks,\n",
    "        hidden_size,\n",
    "        num_heads,\n",
    "        dropout_rate,\n",
    "    ):\n",
    "        super(Transformer, self).__init__()\n",
    "        self.num_blocks = num_blocks\n",
    "        self.hidden_size = hidden_size\n",
    "        self.num_heads = num_heads\n",
    "        self.dropout_rate = dropout_rate\n",
    "\n",
    "        self.embedding_softmax_layer = EmbeddingSharedWeights(vocab_size, hidden_size)\n",
    "        self.encoder_stack = EncoderStack(num_blocks, hidden_size, num_heads, dropout_rate)\n",
    "        self.decoder_stack = DecoderStack(num_blocks, hidden_size, num_heads, dropout_rate)\n",
    "\n",
    "        self.encoder_dropout_layer = nn.Dropout(dropout_rate)\n",
    "        self.decoder_dropout_layer = nn.Dropout(dropout_rate)\n",
    "\n",
    "    def forward(self, encoder_inputs, decoder_inputs):\n",
    "        attention_bias = get_padding_bias(encoder_inputs)\n",
    "        encoder_outputs = self.encode(encoder_inputs, attention_bias)\n",
    "        logits = self.decode(decoder_inputs, encoder_outputs, attention_bias)\n",
    "        return logits\n",
    "\n",
    "    def encode(self, inputs, attention_bias):\n",
    "        embedded_inputs = self.embedding_softmax_layer(inputs)\n",
    "\n",
    "        # add_pos_encoding\n",
    "        length = embedded_inputs.size(1)\n",
    "        pos_encoding = get_position_encoding(length, self.hidden_size, embedded_inputs.device)\n",
    "        encoder_inputs = embedded_inputs + pos_encoding\n",
    "        encoder_inputs = self.encoder_dropout_layer(encoder_inputs)\n",
    "\n",
    "        return self.encoder_stack(encoder_inputs, attention_bias)\n",
    "\n",
    "    def decode(self, inputs, encoder_outputs, attention_bias):\n",
    "        embedded_inputs = self.embedding_softmax_layer(inputs)\n",
    "\n",
    "        # add_pos_encoding\n",
    "        length = embedded_inputs.size(1)\n",
    "        pos_encoding = get_position_encoding(length, self.hidden_size, embedded_inputs.device)\n",
    "        decoder_inputs = embedded_inputs + pos_encoding\n",
    "        decoder_inputs = self.decoder_dropout_layer(decoder_inputs)\n",
    "\n",
    "        decoder_self_attention_bias = get_decoder_self_attention_bias(length, decoder_inputs.device)\n",
    "        decoder_outputs = self.decoder_stack(\n",
    "            decoder_inputs,\n",
    "            encoder_outputs,\n",
    "            decoder_self_attention_bias,\n",
    "            attention_bias\n",
    "        )\n",
    "        logits = self.embedding_softmax_layer(decoder_outputs, mode='linear')\n",
    "\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer = Transformer(\n",
    "    vocab_size,\n",
    "    num_blocks,\n",
    "    num_hidden_size,\n",
    "    num_heads,\n",
    "    dropout_rate,\n",
    ").to(device)\n",
    "\n",
    "optimizer = torch.optim.Adam(transformer.parameters(), lr=0.0002, betas=(0.9, 0.97))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_function(real, pred):\n",
    "    loss = F.cross_entropy(pred.view(-1, pred.size(-1)), real.view(-1), ignore_index=0)\n",
    "    return loss\n",
    "\n",
    "\n",
    "def accuracy_function(real, pred):\n",
    "    predicted_ids = pred.argmax(-1)\n",
    "    correct = (predicted_ids == real).type(pred.dtype)\n",
    "    weights = (real != 0).type(pred.dtype)\n",
    "\n",
    "    return (correct * weights).sum() / weights.sum()\n",
    "\n",
    "\n",
    "def train_step(dataset_inputs):\n",
    "    encoder_inputs, decoder_inputs, decoder_targets = dataset_inputs\n",
    "    encoder_inputs = encoder_inputs.to(device)\n",
    "    decoder_inputs = decoder_inputs.to(device)\n",
    "    decoder_targets = decoder_targets.to(device)\n",
    "\n",
    "    optimizer.zero_grad()\n",
    "    logits = transformer(encoder_inputs, decoder_inputs)\n",
    "    loss = loss_function(decoder_targets, logits)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    accuracy = accuracy_function(decoder_targets, logits)\n",
    "\n",
    "    return loss, accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "template = '{}/{} (epoch {}), Train Loss: {:.4f}, Train Accuracy: {:.4f}, Elapsed Time: {:.2f}'\n",
    "train_losses = []\n",
    "train_accuracies = []\n",
    "start = time.time()\n",
    "for e in range(num_epochs):\n",
    "    transformer.train(True)\n",
    "    for i, batch_train_data in enumerate(iter(data_loader)):\n",
    "        loss, accuracy = train_step(batch_train_data)\n",
    "        train_losses.append(loss.item())\n",
    "        train_accuracies.append(accuracy.item())\n",
    "\n",
    "        if (e * num_batches + i + 1) % 100 == 0:\n",
    "            print(template.format(\n",
    "                e * num_batches + i + 1,\n",
    "                num_epochs * num_batches,\n",
    "                e + 1,\n",
    "                sum(train_losses) / len(train_losses),\n",
    "                sum(train_accuracies) / len(train_accuracies),\n",
    "                time.time() - start\n",
    "            ))\n",
    "\n",
    "            train_losses = []\n",
    "            train_accuracies = []\n",
    "            start = time.time()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "valid_data_num = 10240\n",
    "valid_batch_size = 1024\n",
    "\n",
    "valid_source_sequences, valid_target_sequences = create_dataset(valid_data_num)\n",
    "valid_data_loader = torch.utils.data.DataLoader(\n",
    "    dataset=torch.utils.data.TensorDataset(\n",
    "        tokenize(valid_source_sequences),\n",
    "        tokenize(valid_target_sequences, bos=True),\n",
    "        tokenize(valid_target_sequences, eos=True)\n",
    "    ), batch_size=valid_batch_size, shuffle=False, drop_last=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer.train(False)\n",
    "valid_losses = []\n",
    "valid_accuracies = []\n",
    "\n",
    "for encoder_inputs, decoder_inputs, decoder_targets in iter(valid_data_loader):\n",
    "    encoder_inputs = encoder_inputs.to(device)\n",
    "    decoder_inputs = decoder_inputs.to(device)\n",
    "    decoder_targets = decoder_targets.to(device)\n",
    "\n",
    "    logits = transformer(encoder_inputs, decoder_inputs)\n",
    "    loss = loss_function(decoder_targets, logits)\n",
    "    accuracy = accuracy_function(decoder_targets, logits)\n",
    "\n",
    "    valid_losses.append(loss.item())\n",
    "    valid_accuracies.append(accuracy.item())\n",
    "\n",
    "print('Valid Loss: {:.4f}, Valid Accuracy: {:.4f}'.format(\n",
    "    sum(valid_losses) / len(valid_losses),\n",
    "    sum(valid_accuracies) / len(valid_accuracies),\n",
    "))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}