char_lstm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building an RNN in PyTorch\n",
    "\n",
    "In this notebook, I'll construct a character-level RNN with PyTorch. If you are unfamiliar with character-level RNNs, check out [this great article](http://karpathy.github.io/2015/05/21/rnn-effectiveness/) by Andrej Karpathy. The network will train character by character on some text, then generate new text character by character. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Import all libs\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import Variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('anna.txt', 'r') as f:\n",
    "    \n",
    "    text = f.read()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "83\n"
     ]
    }
   ],
   "source": [
    "charss = set(text)\n",
    "print(len(charss))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have the text, encode it as integers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "chars = tuple(set(text))\n",
    "int2char = dict(enumerate(chars))\n",
    "char2int = {ch: ii for ii, ch in int2char.items()}\n",
    "encoded = np.array([char2int[ch] for ch in text])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('`', 'q', 'z', 'W', 'Q', '6', '8', '\\n', 'E', 'd', 'f', '(', '!', 'x', ' ', '%', 'l', 'L', 'U', '-', ';', 'r', '\"', 'Z', 'p', '.', 'V', 'K', '7', 'X', '&', 'A', 'T', '5', '3', \"'\", 'g', '2', 'Y', 'h', '9', 'y', 't', '/', 'j', ',', 'J', 'N', '0', ':', '@', 'I', 'P', 'S', 'a', 'm', 'w', 'e', 'G', 'v', 'b', 'k', 'C', 'F', 'O', 'D', 'u', '1', 'H', '_', 'o', 'M', '4', 'i', '$', '?', 'B', 's', 'n', '*', 'R', 'c', ')')\n"
     ]
    }
   ],
   "source": [
    "print(chars)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Processing the data\n",
    "\n",
    "We're one-hot encoding the data, so we will make a function to do that.\n",
    "\n",
    "we will also create mini-batches for training. We'll take the encoded characters and split them into multiple sequences, given by `n_seqs` (also refered to as \"batch size\" in other places). Each of those sequences will be `n_steps` long."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def one_hot_encode(arr, n_labels):\n",
    "    \n",
    "    # Initialize the the encoded array\n",
    "    one_hot = np.zeros((np.multiply(*arr.shape), n_labels), dtype=np.float32)\n",
    "    \n",
    "    # Fill the appropriate elements with ones\n",
    "    one_hot[np.arange(one_hot.shape[0]), arr.flatten()] = 1.\n",
    "    \n",
    "    # Finally reshape it to get back to the original array\n",
    "    one_hot = one_hot.reshape((*arr.shape, n_labels))\n",
    "    \n",
    "    return one_hot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_batches(arr, n_seqs, n_steps):\n",
    "    '''Create a generator that returns mini-batches of size\n",
    "       n_seqs x n_steps from arr.\n",
    "    '''\n",
    "    \n",
    "    batch_size = n_seqs * n_steps\n",
    "    n_batches = len(arr)//batch_size\n",
    "    \n",
    "    # Keep only enough characters to make full batches\n",
    "    arr = arr[:n_batches * batch_size]\n",
    "    # Reshape into n_seqs rows\n",
    "    arr = arr.reshape((n_seqs, -1))\n",
    "    \n",
    "    for n in range(0, arr.shape[1], n_steps):\n",
    "        # The features\n",
    "        x = arr[:, n:n+n_steps]\n",
    "        # The targets, shifted by one\n",
    "        y = np.zeros_like(x)\n",
    "        try:\n",
    "            y[:, :-1], y[:, -1] = x[:, 1:], arr[:, n+n_steps]\n",
    "        except IndexError:\n",
    "            y[:, :-1], y[:, -1] = x[:, 1:], arr[:, 0]\n",
    "        yield x, y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining the network with PyTorch\n",
    "\n",
    "Here I'll use PyTorch to define the architecture of the network. We start by defining the layers and operations we want. Then, define a method for the forward pass. I'm also going to write a method for predicting characters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CharRNN(nn.Module):\n",
    "    def __init__(self, tokens, n_steps=100, n_hidden=256, n_layers=2,\n",
    "                               drop_prob=0.5, lr=0.001):\n",
    "        super(CharRNN,self).__init__()\n",
    "        self.drop_prob = drop_prob\n",
    "        self.n_layers = n_layers\n",
    "        self.n_hidden = n_hidden\n",
    "        self.lr = lr\n",
    "        \n",
    "        self.chars = tokens\n",
    "        self.int2char = dict(enumerate(self.chars))\n",
    "        self.char2int = {ch: ii for ii, ch in self.int2char.items()}\n",
    "        \n",
    "        self.dropout = nn.Dropout(drop_prob)\n",
    "        self.lstm = nn.LSTM(len(self.chars), n_hidden, n_layers, \n",
    "                            dropout=drop_prob, batch_first=True)\n",
    "        self.fc = nn.Linear(n_hidden, len(self.chars))\n",
    "        \n",
    "        self.init_weights()\n",
    "        \n",
    "    def forward(self, x, hc):\n",
    "        ''' Forward pass through the network '''\n",
    "        \n",
    "        x, (h, c) = self.lstm(x, hc)\n",
    "        x = self.dropout(x)\n",
    "        \n",
    "        # Stack up LSTM outputs\n",
    "        x = x.view(x.size()[0]*x.size()[1], self.n_hidden)\n",
    "        \n",
    "        x = self.fc(x)\n",
    "        \n",
    "        return x, (h, c)\n",
    "    \n",
    "    def predict(self, char, h=None, cuda=False, top_k=None):\n",
    "        ''' Given a character, predict the next character.\n",
    "        \n",
    "            Returns the predicted character and the hidden state.\n",
    "        '''\n",
    "        if cuda:\n",
    "            self.cuda()\n",
    "        else:\n",
    "            self.cpu()\n",
    "        \n",
    "        if h is None:\n",
    "            h = self.init_hidden(1)\n",
    "        \n",
    "        x = np.array([[self.char2int[char]]])\n",
    "        x = one_hot_encode(x, len(self.chars))\n",
    "        inputs = Variable(torch.from_numpy(x), volatile=True)\n",
    "        if cuda:\n",
    "            inputs = inputs.cuda()\n",
    "        \n",
    "        h = tuple([Variable(each.data, volatile=True) for each in h])\n",
    "        out, h = self.forward(inputs, h)\n",
    "\n",
    "        p = F.softmax(out).data\n",
    "        if cuda:\n",
    "            p = p.cpu()\n",
    "        \n",
    "        if top_k is None:\n",
    "            top_ch = np.arange(len(self.chars))\n",
    "        else:\n",
    "            p, top_ch = p.topk(top_k)\n",
    "            top_ch = top_ch.numpy().squeeze()\n",
    "        \n",
    "        p = p.numpy().squeeze()\n",
    "        char = np.random.choice(top_ch, p=p/p.sum())\n",
    "            \n",
    "        return self.int2char[char], h\n",
    "    \n",
    "    def init_weights(self):\n",
    "        ''' Initialize weights for fully connected layer '''\n",
    "        initrange = 0.1\n",
    "        \n",
    "        # Set bias tensor to all zeros\n",
    "        self.fc.bias.data.fill_(0)\n",
    "        # FC weights as random uniform\n",
    "        self.fc.weight.data.uniform_(-1, 1)\n",
    "        \n",
    "    def init_hidden(self, n_seqs):\n",
    "        ''' Initializes hidden state '''\n",
    "        # Create two new tensors with sizes n_layers x n_seqs x n_hidden,\n",
    "        # initialized to zero, for hidden state and cell state of LSTM\n",
    "        weight = next(self.parameters()).data\n",
    "        return (Variable(weight.new(self.n_layers, n_seqs, self.n_hidden).zero_()),\n",
    "                Variable(weight.new(self.n_layers, n_seqs, self.n_hidden).zero_()))\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(net, data, epochs=10, n_seqs=10, n_steps=50, lr=0.001, clip=5, val_frac=0.1, cuda=False, print_every=10):\n",
    "    ''' Traing a network \n",
    "    \n",
    "        Arguments\n",
    "        ---------\n",
    "        \n",
    "        net: CharRNN network\n",
    "        data: text data to train the network\n",
    "        epochs: Number of epochs to train\n",
    "        n_seqs: Number of mini-sequences per mini-batch, aka batch size\n",
    "        n_steps: Number of character steps per mini-batch\n",
    "        lr: learning rate\n",
    "        clip: gradient clipping\n",
    "        val_frac: Fraction of data to hold out for validation\n",
    "        cuda: Train with CUDA on a GPU\n",
    "        print_every: Number of steps for printing training and validation loss\n",
    "    \n",
    "    '''\n",
    "    \n",
    "    net.train()\n",
    "    opt = torch.optim.Adam(net.parameters(), lr=lr)\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    \n",
    "    # create training and validation data\n",
    "    val_idx = int(len(data)*(1-val_frac))\n",
    "    data, val_data = data[:val_idx], data[val_idx:]\n",
    "    \n",
    "    if cuda:\n",
    "        net.cuda()\n",
    "    \n",
    "    counter = 0\n",
    "    n_chars = len(net.chars)\n",
    "    for e in range(epochs):\n",
    "        h = net.init_hidden(n_seqs)\n",
    "        for x, y in get_batches(data, n_seqs, n_steps):\n",
    "            counter += 1\n",
    "            \n",
    "            # One-hot encode our data and make them Torch tensors\n",
    "            x = one_hot_encode(x, n_chars)\n",
    "            x, y = torch.from_numpy(x), torch.from_numpy(y)\n",
    "            \n",
    "            inputs, targets = Variable(x), Variable(y)\n",
    "            if cuda:\n",
    "                inputs, targets = inputs.cuda(), targets.cuda()\n",
    "\n",
    "            # Creating new variables for the hidden state, otherwise\n",
    "            # we'd backprop through the entire training history\n",
    "            h = tuple([Variable(each.data) for each in h])\n",
    "\n",
    "            net.zero_grad()\n",
    "            \n",
    "            output, h = net.forward(inputs, h)\n",
    "            loss = criterion(output, targets.view(n_seqs*n_steps))\n",
    "\n",
    "            loss.backward()\n",
    "            \n",
    "            # `clip_grad_norm` helps prevent the exploding gradient problem in RNNs / LSTMs.\n",
    "            nn.utils.clip_grad_norm(net.parameters(), clip)\n",
    "\n",
    "            opt.step()\n",
    "            \n",
    "            if counter % print_every == 0:\n",
    "                \n",
    "                # Get validation loss\n",
    "                val_h = net.init_hidden(n_seqs)\n",
    "                val_losses = []\n",
    "                for x, y in get_batches(val_data, n_seqs, n_steps):\n",
    "                    # One-hot encode our data and make them Torch tensors\n",
    "                    x = one_hot_encode(x, n_chars)\n",
    "                    x, y = torch.from_numpy(x), torch.from_numpy(y)\n",
    "                    \n",
    "                    # Creating new variables for the hidden state, otherwise\n",
    "                    # we'd backprop through the entire training history\n",
    "                    val_h = tuple([Variable(each.data, volatile=True) for each in val_h])\n",
    "                    \n",
    "                    inputs, targets = Variable(x, volatile=True), Variable(y, volatile=True)\n",
    "                    if cuda:\n",
    "                        inputs, targets = inputs.cuda(), targets.cuda()\n",
    "\n",
    "                    output, val_h = net.forward(inputs, val_h)\n",
    "                    val_loss = criterion(output, targets.view(n_seqs*n_steps))\n",
    "                \n",
    "                    val_losses.append(val_loss.data[0])\n",
    "                \n",
    "                print(\"Epoch: {}/{}...\".format(e+1, epochs),\n",
    "                      \"Step: {}...\".format(counter),\n",
    "                      \"Loss: {:.4f}...\".format(loss.data[0]),\n",
    "                      \"Val Loss: {:.4f}\".format(np.mean(val_losses)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Time to train\n",
    "\n",
    "Now we can actually train the network. First we'll create the network itself, with some given hyperparameters. Then, define the mini-batches sizes (number of sequences and number of steps), and start the training. With the train function, we can set the number of epochs, the learning rate, and other parameters. Also, we can run the training on a GPU by setting `cuda=True`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "if 'net' in locals():\n",
    "    del net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = CharRNN(chars, n_hidden=512, n_layers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/users/gpu/anupriy/anaconda3/lib/python3.7/site-packages/ipykernel_launcher.py:58: UserWarning: torch.nn.utils.clip_grad_norm is now deprecated in favor of torch.nn.utils.clip_grad_norm_.\n",
      "/users/gpu/anupriy/anaconda3/lib/python3.7/site-packages/ipykernel_launcher.py:74: UserWarning: volatile was removed and now has no effect. Use `with torch.no_grad():` instead.\n",
      "/users/gpu/anupriy/anaconda3/lib/python3.7/site-packages/ipykernel_launcher.py:76: UserWarning: volatile was removed and now has no effect. Use `with torch.no_grad():` instead.\n",
      "/users/gpu/anupriy/anaconda3/lib/python3.7/site-packages/ipykernel_launcher.py:83: UserWarning: invalid index of a 0-dim tensor. This will be an error in PyTorch 0.5. Use tensor.item() to convert a 0-dim tensor to a Python number\n",
      "/users/gpu/anupriy/anaconda3/lib/python3.7/site-packages/ipykernel_launcher.py:87: UserWarning: invalid index of a 0-dim tensor. This will be an error in PyTorch 0.5. Use tensor.item() to convert a 0-dim tensor to a Python number\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 1/2... Step: 10... Loss: 3.3252... Val Loss: 3.3099\n",
      "Epoch: 1/2... Step: 20... Loss: 3.1762... Val Loss: 3.1932\n",
      "Epoch: 1/2... Step: 30... Loss: 3.0831... Val Loss: 3.0728\n",
      "Epoch: 1/2... Step: 40... Loss: 2.8867... Val Loss: 2.9048\n",
      "Epoch: 1/2... Step: 50... Loss: 2.7598... Val Loss: 2.7268\n",
      "Epoch: 1/2... Step: 60... Loss: 2.6001... Val Loss: 2.6202\n",
      "Epoch: 1/2... Step: 70... Loss: 2.5316... Val Loss: 2.5534\n",
      "Epoch: 1/2... Step: 80... Loss: 2.4727... Val Loss: 2.5034\n",
      "Epoch: 1/2... Step: 90... Loss: 2.4452... Val Loss: 2.4590\n",
      "Epoch: 1/2... Step: 100... Loss: 2.3890... Val Loss: 2.4180\n",
      "Epoch: 1/2... Step: 110... Loss: 2.3487... Val Loss: 2.3880\n",
      "Epoch: 1/2... Step: 120... Loss: 2.2861... Val Loss: 2.3587\n",
      "Epoch: 1/2... Step: 130... Loss: 2.3071... Val Loss: 2.3347\n",
      "Epoch: 2/2... Step: 140... Loss: 2.2695... Val Loss: 2.3086\n",
      "Epoch: 2/2... Step: 150... Loss: 2.2391... Val Loss: 2.2867\n",
      "Epoch: 2/2... Step: 160... Loss: 2.2310... Val Loss: 2.2549\n",
      "Epoch: 2/2... Step: 170... Loss: 2.1878... Val Loss: 2.2318\n",
      "Epoch: 2/2... Step: 180... Loss: 2.1282... Val Loss: 2.2327\n",
      "Epoch: 2/2... Step: 190... Loss: 2.0835... Val Loss: 2.1887\n",
      "Epoch: 2/2... Step: 200... Loss: 2.0897... Val Loss: 2.1755\n",
      "Epoch: 2/2... Step: 210... Loss: 2.0927... Val Loss: 2.1472\n",
      "Epoch: 2/2... Step: 220... Loss: 2.0474... Val Loss: 2.1317\n",
      "Epoch: 2/2... Step: 230... Loss: 2.0538... Val Loss: 2.1140\n",
      "Epoch: 2/2... Step: 240... Loss: 2.0352... Val Loss: 2.1008\n",
      "Epoch: 2/2... Step: 250... Loss: 1.9762... Val Loss: 2.0761\n",
      "Epoch: 2/2... Step: 260... Loss: 1.9595... Val Loss: 2.0641\n",
      "Epoch: 2/2... Step: 270... Loss: 1.9765... Val Loss: 2.0491\n"
     ]
    }
   ],
   "source": [
    "n_seqs, n_steps = 128, 100\n",
    "train(net, encoded, epochs=2, n_seqs=n_seqs, n_steps=n_steps, lr=0.001, cuda=True, print_every=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting the best model\n",
    "\n",
    "To set your hyperparameters to get the best performance, you'll want to watch the training and validation losses. If your training loss is much lower than the validation loss, you're overfitting. Increase regularization (more dropout) or use a smaller network. If the training and validation losses are close, you're underfitting so you can increase the size of the network."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After training, we'll save the model so we can load it again later if we need too. Here I'm saving the parameters needed to create the same architecture, the hidden layer hyperparameters and the text characters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = {'n_hidden': net.n_hidden,\n",
    "              'n_layers': net.n_layers,\n",
    "              'state_dict': net.state_dict(),\n",
    "              'tokens': net.chars}\n",
    "with open('rnn.net', 'wb') as f:\n",
    "    torch.save(checkpoint, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sampling\n",
    "\n",
    "Now that the model is trained, we'll want to sample from it. To sample, we pass in a character and have the network predict the next character. Then we take that character, pass it back in, and get another predicted character. Just keep doing this and you'll generate a bunch of text!\n",
    "\n",
    "### Top K sampling\n",
    "\n",
    "Our predictions come from a categorcial probability distribution over all the possible characters. We can make the sampled text more reasonable but less variable by only considering some $K$ most probable characters. This will prevent the network from giving us completely absurd characters while allowing it to introduce some noise and randomness into the sampled text.\n",
    "\n",
    "Typically you'll want to prime the network so you can build up a hidden state. Otherwise the network will start out generating characters at random. In general the first bunch of characters will be a little rough since it hasn't built up a long history of characters to predict from."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample(net, size, prime='The', top_k=None, cuda=False):\n",
    "        \n",
    "    if cuda:\n",
    "        net.cuda()\n",
    "    else:\n",
    "        net.cpu()\n",
    "\n",
    "    net.eval()\n",
    "    \n",
    "    # First off, run through the prime characters\n",
    "    chars = [ch for ch in prime]\n",
    "    h = net.init_hidden(1)\n",
    "    for ch in prime:\n",
    "        char, h = net.predict(ch, h, cuda=cuda, top_k=top_k)\n",
    "\n",
    "    chars.append(char)\n",
    "    \n",
    "    # Now pass in the previous character and get a new one\n",
    "    for ii in range(size):\n",
    "        char, h = net.predict(chars[-1], h, cuda=cuda, top_k=top_k)\n",
    "        chars.append(char)\n",
    "\n",
    "    return ''.join(chars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Annad his\n",
      "with she his the corse the hansiof of the sai\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/users/gpu/anupriy/anaconda3/lib/python3.7/site-packages/ipykernel_launcher.py:49: UserWarning: volatile was removed and now has no effect. Use `with torch.no_grad():` instead.\n",
      "/users/gpu/anupriy/anaconda3/lib/python3.7/site-packages/ipykernel_launcher.py:53: UserWarning: volatile was removed and now has no effect. Use `with torch.no_grad():` instead.\n",
      "/users/gpu/anupriy/anaconda3/lib/python3.7/site-packages/ipykernel_launcher.py:56: UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.\n"
     ]
    }
   ],
   "source": [
    "print(sample(net, 50, prime='Anna', top_k=5, cuda=False))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading a checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('rnn.net', 'rb') as f:\n",
    "    checkpoint = torch.load(f)\n",
    "    \n",
    "loaded = CharRNN(checkpoint['tokens'], n_hidden=checkpoint['n_hidden'], n_layers=checkpoint['n_layers'])\n",
    "loaded.load_state_dict(checkpoint['state_dict'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(sample(loaded, 2000, cuda=True, top_k=5, prime=\"AI\"))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "GPU",
   "language": "python",
   "name": "vqa"
  },
  "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}