Creating extensions using numpy and scipy.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating extensions using numpy and scipy\n",
    "\n",
    "In this notebook, we shall go through two tasks:\n",
    "\n",
    "1. Create a neural network layer with no parameters. \n",
    "  - This calls into **numpy** as part of it's implementation\n",
    "2. Create a neural network layer that has learnable weights\n",
    "  - This calls into **SciPy** as part of it's implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.autograd import Function\n",
    "from torch.autograd import Variable"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameter-less example\n",
    "\n",
    "This layer doesn't particularly do anything useful or mathematically correct.\n",
    "\n",
    "It is aptly named BadFFTFunction\n",
    "\n",
    "**Layer Implementation**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from numpy.fft import rfft2, irfft2\n",
    "class BadFFTFunction(Function):\n",
    "    \n",
    "    def forward(self, input):\n",
    "        numpy_input = input.numpy()\n",
    "        result = abs(rfft2(numpy_input))\n",
    "        return torch.FloatTensor(result)\n",
    "    \n",
    "    def backward(self, grad_output):\n",
    "        numpy_go = grad_output.numpy()\n",
    "        result = irfft2(numpy_go)\n",
    "        return torch.FloatTensor(result)\n",
    "\n",
    "# since this layer does not have any parameters, we can\n",
    "# simply declare this as a function, rather than as an nn.Module class\n",
    "def incorrect_fft(input):\n",
    "    return BadFFTFunction()(input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Example usage of the created layer:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "  6.9997  11.0343   9.7395   6.0724   6.0526\n",
      "  7.0250  11.4841   7.1110   5.6337   8.6441\n",
      "  7.8062  10.9281   9.8279  23.4972   7.4842\n",
      "  6.4962   4.5987   0.7936   3.9360   4.9595\n",
      "  9.7913  10.3732   1.6261   2.0874  14.5295\n",
      "  6.4962   5.7111   1.9793   8.8037   4.9595\n",
      "  7.8062   8.7752   6.4442  14.1250   7.4842\n",
      "  7.0250   5.4642   1.7983   4.4346   8.6441\n",
      "[torch.FloatTensor of size 8x5]\n",
      "\n",
      "Variable containing:\n",
      "-0.0129  0.0330  0.0036 -0.0737  0.2354 -0.0737  0.0036  0.0330\n",
      " 0.0542  0.0986 -0.0382 -0.1137 -0.0944 -0.0973 -0.0172 -0.0021\n",
      "-0.1538 -0.1444  0.0356  0.1590  0.0588 -0.0188 -0.0611  0.0346\n",
      " 0.1511  0.0370 -0.2513 -0.1518  0.1513 -0.2312 -0.0896 -0.1450\n",
      "-0.1668 -0.0814  0.1954  0.1405  0.2191  0.1405  0.1954 -0.0814\n",
      " 0.1511 -0.1450 -0.0896 -0.2312  0.1513 -0.1518 -0.2513  0.0370\n",
      "-0.1538  0.0346 -0.0611 -0.0188  0.0588  0.1590  0.0356 -0.1444\n",
      " 0.0542 -0.0021 -0.0172 -0.0973 -0.0944 -0.1137 -0.0382  0.0986\n",
      "[torch.FloatTensor of size 8x8]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "input = Variable(torch.randn(8, 8), requires_grad=True)\n",
    "result = incorrect_fft(input)\n",
    "print(result.data)\n",
    "result.backward(torch.randn(result.size()))\n",
    "print(input.grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parametrized example\n",
    "\n",
    "This implements a layer with learnable weights.\n",
    "\n",
    "It implements the Cross-correlation with a learnable kernel.\n",
    "\n",
    "In deep learning literature, it's confusingly referred to as Convolution.\n",
    "\n",
    "The backward computes the gradients wrt the input and gradients wrt the filter.\n",
    "\n",
    "**Implementation:**\n",
    "\n",
    "*Please Note that the implementation serves as an illustration, and we did not verify it's correctness*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from scipy.signal import convolve2d, correlate2d\n",
    "from torch.nn.modules.module import Module\n",
    "from torch.nn.parameter import Parameter\n",
    "\n",
    "class ScipyConv2dFunction(Function):\n",
    "    \n",
    "    def forward(self, input, filter):\n",
    "        result = correlate2d(input.numpy(), filter.numpy(), mode='valid')\n",
    "        self.save_for_backward(input, filter)\n",
    "        return torch.FloatTensor(result)\n",
    "    \n",
    "    def backward(self, grad_output):\n",
    "        input, filter = self.saved_tensors\n",
    "        grad_input = convolve2d(grad_output.numpy(), filter.t().numpy(), mode='full')\n",
    "        grad_filter = convolve2d(input.numpy(), grad_output.numpy(), mode='valid')\n",
    "        return torch.FloatTensor(grad_input), torch.FloatTensor(grad_filter)\n",
    "\n",
    "\n",
    "class ScipyConv2d(Module):\n",
    "    \n",
    "    def __init__(self, kh, kw):\n",
    "        super(ScipyConv2d, self).__init__()\n",
    "        self.filter=Parameter(torch.randn(kh, kw))\n",
    "    \n",
    "    def forward(self, input):\n",
    "        return ScipyConv2dFunction()(input, self.filter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Example usage: **"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Parameter containing:\n",
      " 0.0460  0.5052  0.9281\n",
      " 0.8355  1.2642 -0.1283\n",
      " 1.7027 -0.3146 -0.6927\n",
      "[torch.FloatTensor of size 3x3]\n",
      "]\n",
      "Variable containing:\n",
      " 1.4619 -4.0543  0.4391 -0.5423 -4.3719  3.9728 -0.4084 -2.8224\n",
      "-3.6799 -3.9278  4.9932 -3.8952  3.0663  1.6303  2.9775  1.1806\n",
      "-3.1694  2.1434  0.4432  1.6941  1.9344 -0.1196  1.1259  4.3571\n",
      "-0.7934 -1.4610  2.2360  0.6406  0.3729  1.9140  0.2427  0.4298\n",
      "-2.2961 -0.4189  5.6658  0.8090 -1.3030  2.2934  0.7164 -0.0272\n",
      " 1.0649  1.0400 -1.3774 -0.2026 -0.9841  1.7192  3.0843  3.4241\n",
      " 3.2743 -1.8780 -2.3084  0.8508  1.1622  0.6060  2.5559  1.0228\n",
      "-2.3282 -1.1790 -2.4604 -1.9252 -1.3962  1.1054  3.6035  3.1302\n",
      "[torch.FloatTensor of size 8x8]\n",
      "\n",
      "Variable containing:\n",
      " 0.0427  0.7780  1.7383  1.8333  3.8198  0.1135 -3.5576 -4.3994 -0.4354 -0.6021\n",
      " 0.4661  1.2470  2.1080  6.3960  0.6894 -4.5144 -3.2005 -0.2762  0.3508  1.7803\n",
      " 0.8492  0.9083  4.1836  0.6133 -3.4092 -1.8541  0.2254  3.6970  1.0382  0.5031\n",
      " 0.0919  1.7864  1.5422  0.2942  2.0176  1.0741  0.8390  2.6984  2.4786  0.2636\n",
      " 0.2600  0.5248  2.3759  2.1921 -3.4520 -3.2025  2.6008 -0.7395  0.3200  0.0964\n",
      " 0.1632  1.9750  2.5973 -2.0378 -5.2213  1.2097  1.3411  1.6995 -1.4448 -2.6965\n",
      " 0.5332  0.8034 -3.0446 -6.2269 -3.4281 -0.5354 -0.4278 -0.7310 -1.1542  0.7947\n",
      " 0.1243 -1.0476 -2.9011 -5.9247 -2.5209 -3.1030 -4.4343 -2.7956  1.4640  0.0090\n",
      "-0.9033 -0.4323 -2.5873 -1.8884 -1.4657 -1.4747 -0.0032  1.4012 -0.7892 -0.1049\n",
      " 0.0739 -0.7349 -0.3925 -0.9291 -1.1198  0.5321  1.9748  0.1242 -0.4062  0.3108\n",
      "[torch.FloatTensor of size 10x10]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "module = ScipyConv2d(3, 3)\n",
    "print(list(module.parameters()))\n",
    "input = Variable(torch.randn(10, 10), requires_grad=True)\n",
    "output = module(input)\n",
    "print(output)\n",
    "output.backward(torch.randn(8, 8))\n",
    "print(input.grad)"
   ]
  }
 ],
 "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}