LanczosCheck.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Checking the correctness\n",
    "\n",
    "I'm still working on the checks of the Hessian-vector calculations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torchessian import hessian_matmul\n",
    "from torchessian.batch_mode import lanczos\n",
    "from itertools import product\n",
    "from copy import deepcopy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 174,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hessian_matmul_pearl(model, loss_function, v, batch):\n",
    "    \"\"\"\n",
    "        This is a PyTorch implementation of the method introduced\n",
    "        by the article \"Fast Exact Multiplication by the Hessian\"\n",
    "        \n",
    "        - Link: http://www.bcl.hamilton.ie/~barak/papers/nc-hessian.pdf\n",
    "        \n",
    "    \"\"\"\n",
    "    r = 1e-3\n",
    "    model_delta = deepcopy(model)\n",
    "    begin = end = 0\n",
    "    for p in model_delta.parameters():\n",
    "        if not p.requires_grad:\n",
    "            continue\n",
    "        end = begin + p.data.numel()\n",
    "        p_flat = p.data.view(-1)\n",
    "        p_flat += r * v[begin:end]\n",
    "        begin = end\n",
    "\n",
    "\n",
    "    model.zero_grad()\n",
    "\n",
    "    x, y = batch\n",
    "    E = loss_function(model(x), y)\n",
    "    E_delta = loss_function(model_delta(x), y)\n",
    "\n",
    "    E.backward()\n",
    "    E_delta.backward()\n",
    "\n",
    "    grad_w = torch.cat(\n",
    "        list(p.grad.view(1, -1) for p in model.parameters() if p.requires_grad), \n",
    "        1\n",
    "    )\n",
    "    grad_w_delta = torch.cat(\n",
    "        list(p.grad.view(1, -1) for p in model_delta.parameters() if p.requires_grad), \n",
    "        1\n",
    "    )\n",
    "\n",
    "    grad_w.squeeze_()\n",
    "    grad_w_delta.squeeze_()\n",
    "\n",
    "\n",
    "    return (grad_w_delta - grad_w) / r\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 175,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.6698, 0.5066, 0.9960],\n",
      "        [0.5066, 0.6937, 1.0245],\n",
      "        [0.9960, 1.0245, 2.0000]])\n"
     ]
    }
   ],
   "source": [
    "n = 2\n",
    "model = torch.nn.Linear(n, 1)\n",
    "# torch.manual_seed(2019)\n",
    "v = torch.rand(sum(p.data.numel() for p in model.parameters()))\n",
    "x = torch.rand(1000, n)\n",
    "y = x.sum(1)\n",
    "def loss_function(y_hat, y):\n",
    "    return ((y_hat.view(-1) - y.view(-1)) * (y_hat.view(-1) - y.view(-1))).sum() / y.numel()\n",
    "batch = x, y\n",
    "m = 3\n",
    "# Exact Hessian\n",
    "\n",
    "h11 = 2 * (x[:, 0] * x[:, 0]).sum() / x.size(0)\n",
    "h21 = h12 = 2 * (x[:, 0] * x[:, 1]).sum() / x.size(0)\n",
    "h31 = h13 = 2 * x[:, 0].sum() / x.size(0)\n",
    "\n",
    "h22 = 2 * (x[:, 1] * x[:, 1]).sum() / x.size(0)\n",
    "h32 = h23 = 2 * x[:, 1].sum() / x.size(0)\n",
    "\n",
    "h33 = 2\n",
    "\n",
    "H = torch.zeros(3, 3)\n",
    "\n",
    "for i, j in product(range(3), range(3)):\n",
    "    exec(\"H[%d, %d] = h%d%d\" % (i, j, i+1, j+1))\n",
    "    \n",
    "print(H)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 176,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:00<00:00, 841.72it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Batch-mode LANCZOS Algorithm running]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([3.0796, 0.1088, 0.1751])"
      ]
     },
     "execution_count": 176,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.eig(lanczos(model, loss_function, batch, 3, 3)[0])[0][:, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 177,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.0020)"
      ]
     },
     "execution_count": 177,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v = torch.rand(3)\n",
    "h1 = hessian_matmul_pearl(model, loss_function, v, batch)\n",
    "h2 = hessian_matmul(model, loss_function, v, batch)\n",
    "\n",
    "(h1 - h2).norm()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.models import MobileNetV2\n",
    "\n",
    "model = MobileNetV2()\n",
    "torch.manual_seed(2019)\n",
    "x = torch.rand(64, 3, 224, 244)\n",
    "y = torch.randint(2, (64, ))\n",
    "loss_function = torch.nn.CrossEntropyLoss()\n",
    "batch = x, y\n",
    "v = torch.rand(sum(p.data.numel() for p in model.parameters() if p.requires_grad))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 170,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(12318.9453)"
      ]
     },
     "execution_count": 170,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "h1 = hessian_matmul_pearl(model, loss_function, v, batch)\n",
    "h2 = hessian_matmul(model, loss_function, v, batch)\n",
    "\n",
    "(h1 - h2).norm()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-1174.3099, -1073.7307,  -960.3297,  ...,   771.1893,   795.3997,\n",
       "          842.8501])"
      ]
     },
     "execution_count": 171,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "h1.sort()[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 172,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-664.6807, -473.8264, -432.6977,  ...,  315.0118,  340.9702,\n",
       "         385.2315])"
      ]
     },
     "execution_count": 172,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "h2.sort()[0]"
   ]
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}