Test PyTorch grad function

test_pytorch_grad.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.autograd import grad\n",
    "\n",
    "\n",
    "def fwd_gradients(val, x):\n",
    "    dummy = torch.ones_like(val)\n",
    "    dummy.requires_grad = True\n",
    "    g = grad(val, x, grad_outputs=dummy, create_graph=True)[0]\n",
    "    g2 = grad(g, dummy, grad_outputs=torch.ones_like(g), create_graph=True)[0]\n",
    "    print(g.shape, g2.shape)\n",
    "    return g, g2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define Jacobian matrix $J_{i,j}=\\frac{\\partial y_i}{\\partial x_j}$.\n",
    "\n",
    "`torch.autograd.grad(y, x, v)`$:=J^{T} \\cdot v=\\left(\\begin{array}{ccc}{\\frac{\\partial y_{1}}{\\partial x_{1}}} & {\\cdots} & {\\frac{\\partial y_{m}}{\\partial x_{1}}} \\\\ {\\vdots} & {\\ddots} & {\\vdots} \\\\ {\\frac{\\partial y_{1}}{\\partial x_{n}}} & {\\cdots} & {\\frac{\\partial y_{m}}{\\partial x_{n}}}\\end{array}\\right)\\left(\\begin{array}{c}{\\frac{\\partial l}{\\partial y_{1}}} \\\\ {\\vdots} \\\\ {\\frac{\\partial l}{\\partial y_{m}}}\\end{array}\\right)=\\left(\\begin{array}{c}{\\frac{\\partial l}{\\partial x_{1}}} \\\\ {\\vdots} \\\\ {\\frac{\\partial l}{\\partial x_{n}}}\\end{array}\\right)$.\n",
    "\n",
    "Let `dummy = torch.ones_like(y)`.\n",
    "\n",
    "$g_k=$ `grad(y, x, dummy)[k]` $=\\sum_{i=1}^m\\frac{\\partial y_i}{\\partial x_k}$.\n",
    "\n",
    "${g_2}_k=$ `grad(g, dummy, torch.ones_like(g)[k]` $=\\sum_{i=1}^n\\frac{\\partial y_k}{\\partial x_i}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([100, 3, 3])\n",
      "Ground-truth Jacobian summing all y\n",
      "tensor([[ 4.4667,  3.5875,  2.4143],\n",
      "        [-4.0981,  1.4698, -2.3561],\n",
      "        [ 0.1766, -2.9466, -1.2416],\n",
      "        [ 5.2866,  1.2042, -0.8181],\n",
      "        [-0.9223, -0.5854, -1.0042]], grad_fn=<SumBackward1>)\n",
      "\n",
      "Autograd Jacobian summing all y\n",
      "tensor([[ 4.4667,  3.5875,  2.4143],\n",
      "        [-4.0981,  1.4698, -2.3561],\n",
      "        [ 0.1766, -2.9466, -1.2416],\n",
      "        [ 5.2866,  1.2042, -0.8181],\n",
      "        [-0.9223, -0.5854, -1.0042]], grad_fn=<SliceBackward>)\n",
      "\n",
      "Ground-truth Jacobian summing all x\n",
      "tensor([[ 3.0000, -1.1613,  8.6298],\n",
      "        [ 3.0000, -2.6611, -5.3233],\n",
      "        [ 3.0000, -2.2119, -4.7997],\n",
      "        [ 3.0000,  2.9675, -0.2947],\n",
      "        [ 3.0000, -1.3996, -4.1123]], grad_fn=<SumBackward1>)\n",
      "\n",
      "Autograd Jacobian summing all x\n",
      "tensor([[ 3.0000, -1.1613,  8.6298],\n",
      "        [ 3.0000, -2.6611, -5.3233],\n",
      "        [ 3.0000, -2.2119, -4.7997],\n",
      "        [ 3.0000,  2.9675, -0.2947],\n",
      "        [ 3.0000, -1.3996, -4.1123]], grad_fn=<SliceBackward>)\n",
      "\n",
      "tensor(1.1538e-06, grad_fn=<NormBackward0>)\n",
      "tensor(0., grad_fn=<NormBackward0>)\n"
     ]
    }
   ],
   "source": [
    "np.random.seed(42)\n",
    "torch.manual_seed(42)\n",
    "\n",
    "x = torch.zeros(100, 3)\n",
    "x.normal_()\n",
    "x.requires_grad = True\n",
    "y1 = torch.sum(x, dim=-1, keepdim=True)\n",
    "y2 = torch.sin(torch.sum(x, dim=-1, keepdim=True))\n",
    "y3 = torch.sum(x * x, dim=-1, keepdim=True)\n",
    "y = torch.cat((y1, y2, y3), dim=-1)\n",
    "gt_jaco_y1_x = torch.ones_like(x)\n",
    "gt_jaco_y2_x = torch.cos(torch.sum(x, dim=-1, keepdim=True)).repeat(1, 3)\n",
    "gt_jaco_y3_x = 2 * x\n",
    "gt_jaco = torch.stack((gt_jaco_y1_x, gt_jaco_y2_x, gt_jaco_y3_x), dim=1)\n",
    "print(gt_jaco.shape)\n",
    "\n",
    "dummy = torch.ones_like(y)\n",
    "dummy.requires_grad = True\n",
    "g = grad(y, x, grad_outputs=dummy, create_graph=True)[0]\n",
    "g2 = grad(g, dummy, grad_outputs=torch.ones_like(g), create_graph=True)[0]\n",
    "print('Ground-truth Jacobian summing all y')\n",
    "print(torch.sum(gt_jaco[:5], dim=1))\n",
    "print()\n",
    "print('Autograd Jacobian summing all y')\n",
    "print(g[:5])\n",
    "print()\n",
    "print('Ground-truth Jacobian summing all x')\n",
    "print(torch.sum(gt_jaco[:5], dim=2))\n",
    "print()\n",
    "print('Autograd Jacobian summing all x')\n",
    "print(g2[:5])\n",
    "print()\n",
    "print(torch.norm(g - torch.sum(gt_jaco, dim=1)))\n",
    "print(torch.norm(g2 - torch.sum(gt_jaco, dim=2)))"
   ]
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}