PtrNet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch import optim\n",
    "import numpy as np\n",
    "import random\n",
    "from sklearn.model_selection import train_test_split\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def create_argsort(n_samples, seq_len):\n",
    "    \"\"\"argsort task1\"\"\"\n",
    "    data = np.random.randint(seq_len, size=(n_samples, seq_len))\n",
    "    labels = np.argsort(data)\n",
    "    return data, labels\n",
    "\n",
    "def generate_single_seq(length=30, min_len=5, max_len=10):\n",
    "    seq_before = [(random.randint(1, 5)) for x in range(random.randint(min_len, max_len))]\n",
    "    seq_during = [(random.randint(6, 10)) for x in range(random.randint(min_len, max_len))]\n",
    "    seq_after = [random.randint(1, 5) for x in range(random.randint(min_len, max_len))]\n",
    "    seq = seq_before + seq_during + seq_after\n",
    "    seq = seq + ([0] * (length - len(seq)))\n",
    "    return seq, len(seq_before), len(seq_before) + len(seq_during) - 1\n",
    "\n",
    "def generate_set_seq(N, seq_len):\n",
    "    \"\"\"boundary task2\"\"\"\n",
    "    data = []\n",
    "    starts = []\n",
    "    ends = []\n",
    "    for i in range(N):\n",
    "        seq, ind_start, ind_end = generate_single_seq(seq_len)\n",
    "        data.append(seq)\n",
    "        starts.append(ind_start)\n",
    "        ends.append(ind_end)\n",
    "    data = np.array(data)\n",
    "    labels = np.vstack((starts, ends)).T\n",
    "    return data, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "embedding_dim = 100\n",
    "hidden_dim = 100\n",
    "batch_size = 64\n",
    "\n",
    "n_samples = 10000\n",
    "# task 1\n",
    "# input_dim = 10\n",
    "# output_dim = 10\n",
    "# max_trg_len = 10\n",
    "# data, labels = create_argsort(n_samples, 10)\n",
    "\n",
    "# task 2\n",
    "input_dim = 11\n",
    "output_dim = 30\n",
    "max_trg_len = 2\n",
    "data, labels = generate_set_seq(n_samples, 30)\n",
    "X_train, X_test, Y_train, Y_test = train_test_split(data, labels, test_size=0.1, random_state=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Seq2seq(nn.Module):\n",
    "    def __init__(self, input_dim, output_dim, embedding_dim, hidden_dim):\n",
    "        super().__init__()\n",
    "        self.input_dim = input_dim\n",
    "        self.output_dim = output_dim\n",
    "        self.embedding_dim = embedding_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.encoder_embedding = nn.Embedding(input_dim, embedding_dim)\n",
    "        self.decoder_embedding = nn.Embedding(output_dim, embedding_dim)\n",
    "        self.encoder = nn.GRU(embedding_dim, hidden_dim)\n",
    "        self.decoder = nn.GRUCell(embedding_dim, hidden_dim)\n",
    "        self.fc = nn.Linear(embedding_dim + hidden_dim, output_dim)\n",
    "        \n",
    "    def forward(self, inputs, targets):\n",
    "        batch_size = inputs.size(1)\n",
    "        max_len = targets.size(0)\n",
    "        # (L, B)\n",
    "        embedded = self.encoder_embedding(inputs)\n",
    "        targets = self.decoder_embedding(targets)\n",
    "        # (L, B, E)\n",
    "        _, hidden = self.encoder(embedded)\n",
    "        # initialize \n",
    "        decoder_outputs = torch.zeros((max_len, batch_size, self.output_dim)).to(device)\n",
    "        decoder_input = torch.zeros((batch_size, self.embedding_dim)).to(device)\n",
    "        hidden = hidden.squeeze(0) # (B, H)\n",
    "        for i in range(max_len):\n",
    "            hidden = self.decoder(decoder_input, hidden)\n",
    "            # (B, H)\n",
    "            output = F.log_softmax(F.relu(self.fc(torch.cat((decoder_input, hidden), 1))), 1)\n",
    "            decoder_outputs[i] = output\n",
    "            decoder_input = targets[i]\n",
    "\n",
    "        return decoder_outputs\n",
    "    \n",
    "    def predict(self, inputs, max_trg_len):\n",
    "        batch_size = inputs.size(1)\n",
    "        # (L, B)\n",
    "        embedded = self.encoder_embedding(inputs)\n",
    "        # (L, B, E)\n",
    "        _, hidden = self.encoder(embedded)\n",
    "        # initialize \n",
    "        decoder_outputs = torch.zeros((max_trg_len, batch_size, self.output_dim)).to(device)\n",
    "        decoder_input = torch.zeros((batch_size, self.embedding_dim)).to(device)\n",
    "        hidden = hidden.squeeze(0) # (B, H)\n",
    "        for i in range(max_trg_len):\n",
    "            hidden = self.decoder(decoder_input, hidden)\n",
    "            # (B, H)\n",
    "            output = F.log_softmax(F.relu(self.fc(torch.cat((decoder_input, hidden), 1))), 1)\n",
    "            decoder_outputs[i] = output\n",
    "            _, indices = torch.max(output, 1)\n",
    "            decoder_input = self.decoder_embedding(indices)\n",
    "\n",
    "        return decoder_outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def attention(query, key, value):\n",
    "    \"\"\"query[B, H], key[B, L, H], value[B, L, H]\"\"\"\n",
    "    query = query.unsqueeze(1).repeat(1, key.size(1), 1)\n",
    "    # (B, L, H)\n",
    "    score = torch.sum(query * key, -1)\n",
    "    attn = F.softmax(score, -1).unsqueeze(1)\n",
    "    # (B, 1, L)\n",
    "    outputs = torch.matmul(attn, value)\n",
    "    return outputs.squeeze(1), attn.squeeze(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Seq2seqAttn(nn.Module):\n",
    "    def __init__(self, input_dim, output_dim, embedding_dim, hidden_dim):\n",
    "        super().__init__()\n",
    "        self.input_dim = input_dim\n",
    "        self.output_dim = output_dim\n",
    "        self.embedding_dim = embedding_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.encoder_embedding = nn.Embedding(input_dim, embedding_dim)\n",
    "        self.decoder_embedding = nn.Embedding(output_dim, embedding_dim)\n",
    "        self.encoder = nn.GRU(embedding_dim, hidden_dim)\n",
    "        self.decoder = nn.GRUCell(embedding_dim, hidden_dim)\n",
    "        self.fc = nn.Linear(embedding_dim + hidden_dim * 2, output_dim)\n",
    "        \n",
    "    def forward(self, inputs, targets):\n",
    "        batch_size = inputs.size(1)\n",
    "        max_len = targets.size(0)\n",
    "        # (L, B)\n",
    "        embedded = self.encoder_embedding(inputs)\n",
    "        targets = self.decoder_embedding(targets)\n",
    "        # (L, B, E)\n",
    "        encoder_outputs, hidden = self.encoder(embedded)\n",
    "        # initialize \n",
    "        decoder_outputs = torch.zeros((max_len, batch_size, self.output_dim)).to(device)\n",
    "        decoder_input = torch.zeros((batch_size, self.embedding_dim)).to(device)\n",
    "        hidden = hidden.squeeze(0) # (B, H)\n",
    "        for i in range(max_len):\n",
    "            hidden = self.decoder(decoder_input, hidden)\n",
    "            # (B, H)\n",
    "            context, _ = attention(hidden, encoder_outputs.transpose(0, 1), encoder_outputs.transpose(0, 1))\n",
    "            output = F.log_softmax(F.relu(self.fc(torch.cat((decoder_input, hidden, context), 1))), 1)\n",
    "            decoder_outputs[i] = output\n",
    "            decoder_input = targets[i]\n",
    "\n",
    "        return decoder_outputs\n",
    "    \n",
    "    def predict(self, inputs, max_trg_len):\n",
    "        batch_size = inputs.size(1)\n",
    "        # (L, B)\n",
    "        embedded = self.encoder_embedding(inputs)\n",
    "        # (L, B, E)\n",
    "        encoder_outputs, hidden = self.encoder(embedded)\n",
    "        # initialize \n",
    "        decoder_outputs = torch.zeros(max_trg_len, batch_size, self.output_dim).to(device)\n",
    "        decoder_input = torch.zeros((batch_size, self.embedding_dim)).to(device)\n",
    "        hidden = hidden.squeeze(0) # (B, H)\n",
    "        for i in range(max_trg_len):\n",
    "            hidden = self.decoder(decoder_input, hidden)\n",
    "            # (B, H)\n",
    "            context, _ = attention(hidden, encoder_outputs.transpose(0, 1), encoder_outputs.transpose(0, 1))\n",
    "            output = F.log_softmax(F.relu(self.fc(torch.cat((decoder_input,hidden, context), 1))), 1)\n",
    "            decoder_outputs[i] = output\n",
    "            _, indices = torch.max(output, 1)\n",
    "            decoder_input = self.decoder_embedding(indices)\n",
    "\n",
    "        return decoder_outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class PtrNet(nn.Module):\n",
    "    def __init__(self, input_dim, output_dim, embedding_dim, hidden_dim):\n",
    "        super().__init__()\n",
    "        self.input_dim = input_dim\n",
    "        self.output_dim = output_dim\n",
    "        self.embedding_dim = embedding_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "        \n",
    "        self.encoder_embedding = nn.Embedding(input_dim, embedding_dim)\n",
    "        self.decoder_embedding = nn.Embedding(output_dim, embedding_dim)\n",
    "        self.encoder = nn.GRU(embedding_dim, hidden_dim)\n",
    "        self.decoder = nn.GRUCell(embedding_dim, hidden_dim)\n",
    "        \n",
    "        self.W1 = nn.Linear(hidden_dim, hidden_dim, bias=False)\n",
    "        self.W2 = nn.Linear(hidden_dim, hidden_dim, bias=False)\n",
    "        self.v = nn.Linear(hidden_dim, 1, bias=False)\n",
    "    \n",
    "    def forward(self, inputs, targets):\n",
    "        batch_size = inputs.size(1)\n",
    "        max_len = targets.size(0)\n",
    "        # (L, B)\n",
    "        embedded = self.encoder_embedding(inputs)\n",
    "        targets = self.decoder_embedding(targets)\n",
    "        # (L, B, E)\n",
    "        encoder_outputs, hidden = self.encoder(embedded)\n",
    "        # (L, B, H), (1, B, H)\n",
    "        # initialize \n",
    "        decoder_outputs = torch.zeros((max_len, batch_size, self.output_dim)).to(device)\n",
    "        decoder_input = torch.zeros((batch_size, self.embedding_dim)).to(device)\n",
    "        hidden = hidden.squeeze(0) # (B, H)\n",
    "        for i in range(max_len):\n",
    "            hidden = self.decoder(decoder_input, hidden)\n",
    "            # (B, H)\n",
    "            projection1 = self.W1(encoder_outputs)\n",
    "            # (L, B, H)\n",
    "            projection2 = self.W2(hidden)\n",
    "            # (B, H)\n",
    "            output = F.log_softmax(self.v(F.relu(projection1 + projection2)).squeeze(-1).transpose(0, 1), -1)\n",
    "            decoder_outputs[i] = output\n",
    "            decoder_input = targets[i]\n",
    "\n",
    "        return decoder_outputs\n",
    "\n",
    "    \n",
    "    def predict(self, inputs, max_trg_len):\n",
    "        batch_size = inputs.size(1)\n",
    "        # (L, B)\n",
    "        embedded = self.encoder_embedding(inputs)\n",
    "        # (L, B, E)\n",
    "        encoder_outputs, hidden = self.encoder(embedded)\n",
    "        # (L, B, H), (1, B, H)\n",
    "        # initialize \n",
    "        decoder_outputs = torch.zeros(max_trg_len, batch_size, self.output_dim).to(device)\n",
    "        decoder_input = torch.zeros((batch_size, self.embedding_dim)).to(device)\n",
    "        hidden = hidden.squeeze(0) # (B, H)\n",
    "        for i in range(max_trg_len):\n",
    "            hidden = self.decoder(decoder_input, hidden)\n",
    "            # (B, H)\n",
    "            projection1 = self.W1(encoder_outputs)\n",
    "            # (L, B, H)\n",
    "            projection2 = self.W2(hidden)\n",
    "            # (B, H)\n",
    "            a = self.v(F.relu(projection1 + projection2))\n",
    "            output = F.log_softmax(self.v(F.relu(projection1 + projection2)).squeeze(-1).transpose(0, 1), -1)\n",
    "            decoder_outputs[i] = output\n",
    "            _, indices = torch.max(output, 1)\n",
    "            decoder_input = self.decoder_embedding(indices)\n",
    "\n",
    "        return decoder_outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def train(model, X, Y, batch_size, n_epochs):\n",
    "    model.train()\n",
    "    optimizer = optim.Adam(model.parameters())\n",
    "    criterion = nn.NLLLoss().to(device)\n",
    "    n_samples = X.shape[0]\n",
    "    seq_len = X.shape[1]\n",
    "    for epoch in range(n_epochs + 1):\n",
    "        epoch_loss = 0\n",
    "        for i in range(0, n_samples, batch_size):\n",
    "            batch_X = X[i:i + batch_size]\n",
    "            batch_Y = Y[i:i + batch_size]\n",
    "            batch_X = torch.LongTensor(batch_X).transpose(0,1).to(device)\n",
    "            batch_Y = torch.LongTensor(batch_Y).transpose(0,1).to(device)\n",
    "            outputs = model(batch_X, batch_Y).view(-1, seq_len)\n",
    "            targets = batch_Y.view(-1)\n",
    "            \n",
    "            loss = criterion(outputs, targets)\n",
    "            epoch_loss += loss.item()\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "        print('epoch: {} | total loss: {:.4f}'.format(epoch, epoch_loss))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def rand_evaluate(model):\n",
    "    \"\"\"for task 1\"\"\"\n",
    "    data = np.random.randint(seq_len, size=(1, max_trg_len))\n",
    "    print('>', data.flatten())\n",
    "    print('=', np.argsort(data).flatten())\n",
    "    inputs = torch.from_numpy(data.reshape(seq_len, -1)).to(device)\n",
    "    outputs = model.predict(inputs, max_trg_len).squeeze(1)\n",
    "    print('<', torch.max(outputs, 1)[1].cpu().numpy())\n",
    "\n",
    "def accuracy(model, X, Y):\n",
    "    batch_size = X.shape[0]\n",
    "    inputs = torch.LongTensor(X).transpose(0,1).to(device)\n",
    "    probs = model.predict(inputs, max_trg_len)\n",
    "    _, indices = torch.max(probs, 2)\n",
    "    predicted = indices.t().cpu().numpy()\n",
    "    correct_count = sum([1 if all(p==y) else 0 for p, y in zip(predicted, Y)])\n",
    "    print('Acc: {:.2f}% ({}/{})'.format(correct_count / batch_size * 100, correct_count, batch_size))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 | total loss: 86.9745\n",
      "epoch: 1 | total loss: 0.3416\n",
      "epoch: 2 | total loss: 0.0915\n",
      "epoch: 3 | total loss: 0.0412\n",
      "epoch: 4 | total loss: 0.0231\n",
      "epoch: 5 | total loss: 0.0147\n",
      "epoch: 6 | total loss: 0.0101\n",
      "epoch: 7 | total loss: 0.0073\n",
      "epoch: 8 | total loss: 0.0055\n",
      "epoch: 9 | total loss: 0.0043\n",
      "epoch: 10 | total loss: 0.0034\n",
      "epoch: 11 | total loss: 0.0028\n",
      "epoch: 12 | total loss: 0.0023\n",
      "epoch: 13 | total loss: 0.0019\n",
      "epoch: 14 | total loss: 0.0016\n",
      "epoch: 15 | total loss: 0.0014\n",
      "epoch: 16 | total loss: 0.0012\n",
      "epoch: 17 | total loss: 0.0010\n",
      "epoch: 18 | total loss: 0.0009\n",
      "epoch: 19 | total loss: 0.0008\n",
      "epoch: 20 | total loss: 0.0007\n"
     ]
    }
   ],
   "source": [
    "# model = Seq2seq(input_dim, output_dim, embedding_dim, hidden_dim).to(device)\n",
    "# model = Seq2seqAttn(input_dim, output_dim, embedding_dim, hidden_dim).to(device)\n",
    "model = PtrNet(input_dim, output_dim, embedding_dim, hidden_dim).to(device)\n",
    "train(model, X_train, Y_train, batch_size, 20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Acc: 100.00% (1000/1000)\n"
     ]
    }
   ],
   "source": [
    "accuracy(model, X_test, Y_test)"
   ]
  }
 ],
 "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.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}