4-2.Seq2Seq(Attention)/Seq2Seq(Attention).ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# code by Tae Hwan Jung @graykode\n",
        "# Reference : https://github.com/hunkim/PyTorchZeroToAll/blob/master/14_2_seq2seq_att.py\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "# S: Symbol that shows starting of decoding input\n",
        "# E: Symbol that shows starting of decoding output\n",
        "# P: Symbol that will fill in blank sequence if current batch data size is short than time steps\n",
        "\n",
        "def make_batch():\n",
        "    input_batch = [np.eye(n_class)[[word_dict[n] for n in sentences[0].split()]]]\n",
        "    output_batch = [np.eye(n_class)[[word_dict[n] for n in sentences[1].split()]]]\n",
        "    target_batch = [[word_dict[n] for n in sentences[2].split()]]\n",
        "\n",
        "    # make tensor\n",
        "    return torch.FloatTensor(input_batch), torch.FloatTensor(output_batch), torch.LongTensor(target_batch)\n",
        "\n",
        "class Attention(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Attention, self).__init__()\n",
        "        self.enc_cell = nn.RNN(input_size=n_class, hidden_size=n_hidden, dropout=0.5)\n",
        "        self.dec_cell = nn.RNN(input_size=n_class, hidden_size=n_hidden, dropout=0.5)\n",
        "\n",
        "        # Linear for attention\n",
        "        self.attn = nn.Linear(n_hidden, n_hidden)\n",
        "        self.out = nn.Linear(n_hidden * 2, n_class)\n",
        "\n",
        "    def forward(self, enc_inputs, hidden, dec_inputs):\n",
        "        enc_inputs = enc_inputs.transpose(0, 1)  # enc_inputs: [n_step(=n_step, time step), batch_size, n_class]\n",
        "        dec_inputs = dec_inputs.transpose(0, 1)  # dec_inputs: [n_step(=n_step, time step), batch_size, n_class]\n",
        "\n",
        "        # enc_outputs : [n_step, batch_size, num_directions(=1) * n_hidden], matrix F\n",
        "        # enc_hidden : [num_layers(=1) * num_directions(=1), batch_size, n_hidden]\n",
        "        enc_outputs, enc_hidden = self.enc_cell(enc_inputs, hidden)\n",
        "\n",
        "        trained_attn = []\n",
        "        hidden = enc_hidden\n",
        "        n_step = len(dec_inputs)\n",
        "        model = torch.empty([n_step, 1, n_class])\n",
        "\n",
        "        for i in range(n_step):  # each time step\n",
        "            # dec_output : [n_step(=1), batch_size(=1), num_directions(=1) * n_hidden]\n",
        "            # hidden : [num_layers(=1) * num_directions(=1), batch_size(=1), n_hidden]\n",
        "            dec_output, hidden = self.dec_cell(dec_inputs[i].unsqueeze(0), hidden)\n",
        "            attn_weights = self.get_att_weight(dec_output, enc_outputs)  # attn_weights : [1, 1, n_step]\n",
        "            trained_attn.append(attn_weights.squeeze().data.numpy())\n",
        "\n",
        "            # matrix-matrix product of matrices [1,1,n_step] x [1,n_step,n_hidden] = [1,1,n_hidden]\n",
        "            context = attn_weights.bmm(enc_outputs.transpose(0, 1))\n",
        "            dec_output = dec_output.squeeze(0)  # dec_output : [batch_size(=1), num_directions(=1) * n_hidden]\n",
        "            context = context.squeeze(1)  # [1, num_directions(=1) * n_hidden]\n",
        "            model[i] = self.out(torch.cat((dec_output, context), 1))\n",
        "\n",
        "        # make model shape [n_step, n_class]\n",
        "        return model.transpose(0, 1).squeeze(0), trained_attn\n",
        "\n",
        "    def get_att_weight(self, dec_output, enc_outputs):  # get attention weight one 'dec_output' with 'enc_outputs'\n",
        "        n_step = len(enc_outputs)\n",
        "        attn_scores = torch.zeros(n_step)  # attn_scores : [n_step]\n",
        "\n",
        "        for i in range(n_step):\n",
        "            attn_scores[i] = self.get_att_score(dec_output, enc_outputs[i])\n",
        "\n",
        "        # Normalize scores to weights in range 0 to 1\n",
        "        return F.softmax(attn_scores).view(1, 1, -1)\n",
        "\n",
        "    def get_att_score(self, dec_output, enc_output):  # enc_outputs [batch_size, num_directions(=1) * n_hidden]\n",
        "        score = self.attn(enc_output)  # score : [batch_size, n_hidden]\n",
        "        return torch.dot(dec_output.view(-1), score.view(-1))  # inner product make scalar value\n",
        "\n",
        "if __name__ == '__main__':\n",
        "    n_step = 5 # number of cells(= number of Step)\n",
        "    n_hidden = 128 # number of hidden units in one cell\n",
        "\n",
        "    sentences = ['ich mochte ein bier P', 'S i want a beer', 'i want a beer E']\n",
        "\n",
        "    word_list = \" \".join(sentences).split()\n",
        "    word_list = list(set(word_list))\n",
        "    word_dict = {w: i for i, w in enumerate(word_list)}\n",
        "    number_dict = {i: w for i, w in enumerate(word_list)}\n",
        "    n_class = len(word_dict)  # vocab list\n",
        "\n",
        "    # hidden : [num_layers(=1) * num_directions(=1), batch_size, n_hidden]\n",
        "    hidden = torch.zeros(1, 1, n_hidden)\n",
        "\n",
        "    model = Attention()\n",
        "    criterion = nn.CrossEntropyLoss()\n",
        "    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
        "\n",
        "    input_batch, output_batch, target_batch = make_batch()\n",
        "\n",
        "    # Train\n",
        "    for epoch in range(2000):\n",
        "        optimizer.zero_grad()\n",
        "        output, _ = model(input_batch, hidden, output_batch)\n",
        "\n",
        "        loss = criterion(output, target_batch.squeeze(0))\n",
        "        if (epoch + 1) % 400 == 0:\n",
        "            print('Epoch:', '%04d' % (epoch + 1), 'cost =', '{:.6f}'.format(loss))\n",
        "\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "    # Test\n",
        "    test_batch = [np.eye(n_class)[[word_dict[n] for n in 'SPPPP']]]\n",
        "    test_batch = torch.FloatTensor(test_batch)\n",
        "    predict, trained_attn = model(input_batch, hidden, test_batch)\n",
        "    predict = predict.data.max(1, keepdim=True)[1]\n",
        "    print(sentences[0], '->', [number_dict[n.item()] for n in predict.squeeze()])\n",
        "\n",
        "    # Show Attention\n",
        "    fig = plt.figure(figsize=(5, 5))\n",
        "    ax = fig.add_subplot(1, 1, 1)\n",
        "    ax.matshow(trained_attn, cmap='viridis')\n",
        "    ax.set_xticklabels([''] + sentences[0].split(), fontdict={'fontsize': 14})\n",
        "    ax.set_yticklabels([''] + sentences[2].split(), fontdict={'fontsize': 14})\n",
        "    plt.show()"
      ],
      "outputs": [],
      "execution_count": null
    }
  ],
  "metadata": {
    "anaconda-cloud": {},
    "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.6.1"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 4
}