example-sudoku.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# OptNet/qpth Example Sudoku Notebook\n",
    "\n",
    "*By [Brandon Amos](https://bamos.github.io) and [J. Zico Kolter](http://zicokolter.com/).*\n",
    "\n",
    "---\n",
    "\n",
    "This notebook is released along with our paper\n",
    "[OptNet: Differentiable Optimization as a Layer in Neural Networks](https://arxiv.org/abs/1703.00443).\n",
    "\n",
    "This notebook shows an example of constructing an\n",
    "OptNet layer in PyTorch with our [qpth library](https://github.com/locuslab/qpth)\n",
    "to solve [the game Sudoku](https://en.wikipedia.org/wiki/Sudoku)\n",
    "as a prediction problem from data.\n",
    "See [our qpth documentation page](https://locuslab.github.io/qpth/)\n",
    "for more details on how to use `qpth`.\n",
    "The experiments for our paper that use this library are in\n",
    "[this repo](https://github.com/locuslab/optnet).\n",
    "Specifically [here](https://github.com/locuslab/optnet/tree/master/sudoku)\n",
    "is the full source code for the publihsed version of Sudoku.\n",
    "\n",
    "\n",
    "## Setup and Dependencies\n",
    "\n",
    "+ Python/numpy/[PyTorch](https://pytorch.org)\n",
    "+ [qpth](https://github.com/locuslab/qpth):\n",
    "  *Our fast QP solver for PyTorch released in conjunction with this paper.*\n",
    "+ [bamos/block](https://github.com/bamos/block):\n",
    "  *Our intelligent block matrix library for numpy, PyTorch, and beyond.*\n",
    "+ Optional: [bamos/setGPU](https://github.com/bamos/setGPU):\n",
    "  A small library to set `CUDA_VISIBLE_DEVICES` on multi-GPU systems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import sys\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "import torch.nn as nn\n",
    "from torch.autograd import Function, Variable\n",
    "from torch.nn.parameter import Parameter\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from qpth.qp import QPFunction\n",
    "\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "plt.style.use('bmh')\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup: Download the data and pretrained model\n",
    "\n",
    "+ The pre-trained model is for later.\n",
    "  The following command should download everything to a tmp directory for you\n",
    "  if you have the `wget` and `tar` commands installed.\n",
    "+ (Sorry for the bad form here)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "tmpDir = \"/tmp/optnet.sudoku\"\n",
    "cmd = ('mkdir {}; cd {} &&'\n",
    "       'wget \"http://joule.isr.cs.cmu.edu:11235/optnet/arxiv.v1.sudoku.tgz\" && '\n",
    "       'tar xf arxiv.v1.sudoku.tgz').format(*[tmpDir]*2)\n",
    "dataDir = os.path.join(tmpDir, 'arxiv.v1.sudoku')\n",
    "assert os.system(cmd) == 0\n",
    "\n",
    "sys.path.append(tmpDir+'/arxiv.v1.sudoku')\n",
    "import models # From /tmp/optnet.sudoku/arxiv.v1.sudoku/models.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "testPct = 0.1\n",
    "\n",
    "with open('{}/2/features.pt'.format(dataDir), 'rb') as f:\n",
    "    X = torch.load(f)\n",
    "with open('{}/2/labels.pt'.format(dataDir), 'rb') as f:\n",
    "    Y = torch.load(f)\n",
    "\n",
    "N, nFeatures = X.size(0), int(np.prod(X.size()[1:]))\n",
    "\n",
    "nTrain = int(N*(1.-testPct))\n",
    "nTest = N-nTrain\n",
    "\n",
    "trainX = X[:nTrain]\n",
    "trainY = Y[:nTrain]\n",
    "testX = X[nTrain:]\n",
    "testY = Y[nTrain:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What the data for the Sudoku task looks like\n",
    "\n",
    "The inputs are incomplete boards and the outputs\n",
    "are the completed boards. Here's what the first\n",
    "input and output in the test set looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First testing example input (unsolved Sudoku board):  \n",
      " 0  0  2  4\n",
      " 3  0  0  1\n",
      " 0  4  0  0\n",
      " 0  0  3  0\n",
      "[torch.LongTensor of size 4x4]\n",
      "\n",
      "First testing example output (solved Sudoku board):  \n",
      " 1  3  2  4\n",
      " 3  2  4  1\n",
      " 2  4  1  3\n",
      " 4  1  3  2\n",
      "[torch.LongTensor of size 4x4]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def decode_onehot(encoded_board):\n",
    "    \"\"\"Take the unique argmax of the one-hot encoded board.\"\"\"\n",
    "    v,I = torch.max(encoded_board, 0)\n",
    "    return ((v>0).long()*(I+1)).squeeze()\n",
    "\n",
    "print(\"First testing example input (unsolved Sudoku board): \", decode_onehot(testX[0]))\n",
    "print(\"First testing example output (solved Sudoku board): \", decode_onehot(testY[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You may have noticed that we had to decode those examples.\n",
    "That's because they're actually *one-hot encoded* for how\n",
    "we're going to model the task.\n",
    "That means that instead of representing the values as \n",
    "something between 1 and 4, they're represented\n",
    "as a 4-dimensional vector with a 1 in the index of the value.\n",
    "Here's what the same first example from the test set\n",
    "actually looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First test example input one-hot encoded (unsolved Sudoku board):  \n",
      "(0 ,.,.) = \n",
      "  0  0  0  0\n",
      "  0  0  0  1\n",
      "  0  0  0  0\n",
      "  0  0  0  0\n",
      "\n",
      "(1 ,.,.) = \n",
      "  0  0  1  0\n",
      "  0  0  0  0\n",
      "  0  0  0  0\n",
      "  0  0  0  0\n",
      "\n",
      "(2 ,.,.) = \n",
      "  0  0  0  0\n",
      "  1  0  0  0\n",
      "  0  0  0  0\n",
      "  0  0  1  0\n",
      "\n",
      "(3 ,.,.) = \n",
      "  0  0  0  1\n",
      "  0  0  0  0\n",
      "  0  1  0  0\n",
      "  0  0  0  0\n",
      "[torch.FloatTensor of size 4x4x4]\n",
      "\n",
      "First test example output one-hot encoded (solved Sudoku board):  \n",
      "(0 ,.,.) = \n",
      "  1  0  0  0\n",
      "  0  0  0  1\n",
      "  0  0  1  0\n",
      "  0  1  0  0\n",
      "\n",
      "(1 ,.,.) = \n",
      "  0  0  1  0\n",
      "  0  1  0  0\n",
      "  1  0  0  0\n",
      "  0  0  0  1\n",
      "\n",
      "(2 ,.,.) = \n",
      "  0  1  0  0\n",
      "  1  0  0  0\n",
      "  0  0  0  1\n",
      "  0  0  1  0\n",
      "\n",
      "(3 ,.,.) = \n",
      "  0  0  0  1\n",
      "  0  0  1  0\n",
      "  0  1  0  0\n",
      "  1  0  0  0\n",
      "[torch.FloatTensor of size 4x4x4]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"First test example input one-hot encoded (unsolved Sudoku board): \", testX[0])\n",
    "print(\"First test example output one-hot encoded (solved Sudoku board): \", testY[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Defining a model for this task\n",
    "\n",
    "We've now turned (mini-)Sudoku into a machine learning task that\n",
    "you can apply any model and learning algorithm to. \n",
    "In this notebook, we'll just show how to initialize and train\n",
    "an OptNet model for this task.\n",
    "However you can play around and swap this out for any\n",
    "model you want!\n",
    "Check out [our baseline models](https://github.com/locuslab/optnet/blob/master/sudoku/models.py)\n",
    "if you're interested.\n",
    "\n",
    "Sudoku is actually an integer programming problem but\n",
    "we can relax it to an LP (or LP with a small ridge term,\n",
    "which we'll actually use) that can be expressed as:\n",
    "\n",
    "```\n",
    "y* = argmin_y 0.5 eps y^T y - p^T y\n",
    "         s.t. Ay  = b\n",
    "               y >= 0\n",
    "```\n",
    "\n",
    "To quickly explain this, the quadratic term `0.5 eps y^T y`\n",
    "is a small ridge term so we can use `qpth`,\n",
    "`p` is the (flattened) one-hot encoded input,\n",
    "the `-p^T y` term constrains the solution to contain\n",
    "the same pieces as the unsolved board,\n",
    "and the linear equality constraints `Ay = b`\n",
    "encode the constraints of Sudoku (the row, columns,\n",
    "and sub-blocks must contain all of the digits).\n",
    "\n",
    "If you want to check your understanding of this:\n",
    "\n",
    "1. What do some example constraints `a_i^T y = b_i` look like?\n",
    "2. What happens if we remove the linear equality constraint?\n",
    "\n",
    "Implementing this model is just a few lines of PyTorch with our qpth library.\n",
    "Note that in this notebook we'll just execute this on the CPU,\n",
    "but for performance reasons you should use a GPU for serious\n",
    "experiments:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class OptNet(nn.Module):\n",
    "    def __init__(self, n, Qpenalty):\n",
    "        super().__init__()\n",
    "        nx = (n**2)**3\n",
    "        self.Q = Variable(Qpenalty*torch.eye(nx).double())\n",
    "        self.G = Variable(-torch.eye(nx).double())\n",
    "        self.h = Variable(torch.zeros(nx).double())\n",
    "        A_shape = (40, 64) # Somewhat magic, it's from the true solution.\n",
    "        self.A = Parameter(torch.rand(A_shape).double())\n",
    "        self.b = Variable(torch.ones(A_shape[0]).double())\n",
    "\n",
    "    def forward(self, puzzles):\n",
    "        nBatch = puzzles.size(0)\n",
    "\n",
    "        p = -puzzles.view(nBatch, -1)\n",
    "\n",
    "        return QPFunction(verbose=-1)(\n",
    "            self.Q, p.double(), self.G, self.h, self.A, self.b\n",
    "        ).float().view_as(puzzles)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's it! Let's randomly initialize this model and see what it does on the first test set example. What do you expect?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First test example input (unsolved Sudoku board):  \n",
      " 0  0  2  4\n",
      " 3  0  0  1\n",
      " 0  4  0  0\n",
      " 0  0  3  0\n",
      "[torch.LongTensor of size 4x4]\n",
      "\n",
      "First test example output (TRUE solved Sudoku board):  \n",
      " 1  3  2  4\n",
      " 3  2  4  1\n",
      " 2  4  1  3\n",
      " 4  1  3  2\n",
      "[torch.LongTensor of size 4x4]\n",
      "\n",
      "First test example prediction:  \n",
      " 3  4  1  2\n",
      " 2  3  3  2\n",
      " 3  1  2  3\n",
      " 1  1  2  2\n",
      "[torch.LongTensor of size 4x4]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "model = OptNet(2, 0.1)\n",
    "pred = model(Variable(testX[0].unsqueeze(0))).squeeze().data\n",
    "\n",
    "print(\"First test example input (unsolved Sudoku board): \", decode_onehot(testX[0]))\n",
    "print(\"First test example output (TRUE solved Sudoku board): \", decode_onehot(testY[0]))\n",
    "print(\"First test example prediction: \", decode_onehot(pred))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Wow that prediction is way off!! That's expected since the model was randomly initialized. Note that at this point, some of the constraints actually make it impossible to match the unsolved board (like the `4` at the top right corner)\n",
    "\n",
    "Let's look a random nonsense constraint that the model just satisfied. Here are the coefficients in the first row, `a_1` and `b_1`. The last line here\n",
    "shows that the constraint is acutally satisfied (up to machine precision)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First row of A:\n",
      " [ 0.49  0.74  0.66  0.42  0.95  0.73  0.83  0.2   0.39  0.63  0.36  0.25\n",
      "  0.4   0.26  0.81  0.88  0.98  0.61  0.89  0.9   0.51  0.15  0.11  0.82\n",
      "  0.14  0.22  0.56  0.13  0.64  0.82  0.92  0.25  0.2   0.04  0.25  0.63\n",
      "  0.07  0.68  0.78  0.34  0.62  0.01  0.72  0.54  0.07  0.41  0.43  0.18\n",
      "  0.02  0.21  0.62  0.81  0.3   0.97  0.29  0.51  0.87  0.43  0.6   0.14\n",
      "  0.15  0.16  0.15  0.69]\n",
      "------------------------------\n",
      "First entry of b:  1.0\n",
      "------------------------------\n",
      "a0^T z - b:  -5.92086668583e-09\n"
     ]
    }
   ],
   "source": [
    "np.set_printoptions(precision=2)\n",
    "a0 = model.A[0].data.numpy()\n",
    "b0 = model.b.data[0]\n",
    "z = pred.numpy().ravel()\n",
    "print('First row of A:\\n', a0)\n",
    "print('-'*30)\n",
    "print('First entry of b: ', b0)\n",
    "print('-'*30)\n",
    "print('a0^T z - b: ', np.dot(a0, z) - b0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training the model\n",
    "\n",
    "Let's start training this model my comparing the predictions\n",
    "to the true solutions and taking gradient steps.\n",
    "This takes a while to run (overnight on a GPU), so here\n",
    "we'll just take 10 steps through the first 10 training examples\n",
    "to illustrate what the full training would look like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iteration 0, loss = 0.24\n",
      "Iteration 1, loss = 0.24\n",
      "Iteration 2, loss = 0.24\n",
      "Iteration 3, loss = 0.24\n",
      "Iteration 4, loss = 0.24\n",
      "Iteration 5, loss = 0.23\n",
      "Iteration 6, loss = 0.24\n",
      "Iteration 7, loss = 0.24\n",
      "Iteration 8, loss = 0.22\n",
      "Iteration 9, loss = 0.24\n"
     ]
    }
   ],
   "source": [
    "loss_fn = torch.nn.MSELoss()\n",
    "\n",
    "# Initialize the optimizer.\n",
    "learning_rate = 1e-3\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
    "\n",
    "for t in range(10):\n",
    "    x_batch = Variable(trainX[t].unsqueeze(0))\n",
    "    y_batch = Variable(trainY[t].unsqueeze(0))\n",
    "    \n",
    "    # Forward pass: compute predicted y by passing x to the model.\n",
    "    y_pred = model(x_batch)\n",
    "\n",
    "    # Compute and print loss.\n",
    "    loss = loss_fn(y_pred, y_batch)\n",
    "    print('Iteration {}, loss = {:.2f}'.format(t, loss.data[0]))\n",
    "\n",
    "    # Before the backward pass, use the optimizer object to zero all of the\n",
    "    # gradients for the variables it will update (which are the learnable weights\n",
    "    # of the model)\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    # Backward pass: compute gradient of the loss with respect to model\n",
    "    # parameters\n",
    "    loss.backward()\n",
    "\n",
    "    # Calling the step function on an Optimizer makes an update to its\n",
    "    # parameters\n",
    "    optimizer.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Looking at a pre-trained model\n",
    "\n",
    "Imagine you kept that running for a while.\n",
    "Let's load my pre-trained model we downloaded earlier and\n",
    "see the predictions on the first test example again:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First test example input (unsolved Sudoku board):  \n",
      " 0  0  2  4\n",
      " 3  0  0  1\n",
      " 0  4  0  0\n",
      " 0  0  3  0\n",
      "[torch.LongTensor of size 4x4]\n",
      "\n",
      "First test example output (TRUE solved Sudoku board):  \n",
      " 1  3  2  4\n",
      " 3  2  4  1\n",
      " 2  4  1  3\n",
      " 4  1  3  2\n",
      "[torch.LongTensor of size 4x4]\n",
      "\n",
      "First test example prediction:  \n",
      " 1  3  2  4\n",
      " 3  2  4  1\n",
      " 2  4  1  3\n",
      " 4  1  3  2\n",
      "[torch.LongTensor of size 4x4]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "A_file = os.path.join(tmpDir, 'arxiv.v1.sudoku', 'pretrained-optnet-A.pth')\n",
    "trainedA = torch.load(A_file)\n",
    "\n",
    "trainedModel = OptNet(2, 0.2)\n",
    "trainedModel.A.data = trainedA\n",
    "\n",
    "pred = trainedModel(Variable(testX[0].unsqueeze(0))).data.squeeze()\n",
    "\n",
    "print(\"First test example input (unsolved Sudoku board): \", decode_onehot(testX[0]))\n",
    "print(\"First test example output (TRUE solved Sudoku board): \", decode_onehot(testY[0]))\n",
    "print(\"First test example prediction: \", decode_onehot(pred))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We did it! With just a few lines of code we've trained\n",
    "an intuitive model that solves Sudoku.\n",
    "\n",
    "As a closing note, what does the trained `A` matrix look like?\n",
    "With this formulation, we don't expect it to be the nice,\n",
    "sparse coefficient matrix encoding the rules we typically\n",
    "think of Sudoku as since any row-transformed\n",
    "version of this matrix is an equivalent valid solution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\n",
       " 1.9467  0.9184 -0.3419  ...  -3.1399 -0.6383  1.9654\n",
       " 1.2744 -3.1572  0.0089  ...   5.6140  0.4215 -1.9954\n",
       " 2.1541  0.9581  0.6416  ...   1.2639 -1.1828  3.2781\n",
       "          ...             ⋱             ...          \n",
       "-2.5137  1.8118  0.8839  ...   1.8048  0.6488  1.2468\n",
       "-1.2608 -1.9652  1.1355  ...   3.6201  0.9974  9.0308\n",
       "-4.3232 -0.6153  0.8325  ...   1.4742  2.2705 -1.1042\n",
       "[torch.DoubleTensor of size 40x64]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainedA"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}