PytorchHook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. pytorch hooks  -> x.retain_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0., 2., 4., 6.], grad_fn=<AddBackward0>)\n",
      "tensor(28., grad_fn=<DotBackward>)\n",
      "x.requires_grad: True\n",
      "y.requires_grad: True\n",
      "z.requires_grad: True\n",
      "w.requires_grad: True\n",
      "o.requires_grad: True\n",
      "x.grad: tensor([0., 1., 2., 3.])\n",
      "y.grad: tensor([0., 1., 2., 3.])\n",
      "w.grad: tensor([0., 2., 4., 6.])\n",
      "z.grad: tensor([0., 1., 2., 3.])\n",
      "o.grad: tensor(1.)\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "x = torch.Tensor([0,1,2,3]).requires_grad_()\n",
    "y = torch.Tensor([0,1,2,3]).requires_grad_()\n",
    "w = torch.Tensor([0,1,2,3]).requires_grad_()\n",
    "\n",
    "z = x + y\n",
    "z.retain_grad()  # if we save grad\n",
    "output = w.matmul(z)\n",
    "output.retain_grad() #if we save grad\n",
    "output.backward()\n",
    "\n",
    "print(z)\n",
    "print(output)\n",
    "print('x.requires_grad:', x.requires_grad) # True\n",
    "print('y.requires_grad:', y.requires_grad) # True\n",
    "print('z.requires_grad:', z.requires_grad) # True\n",
    "print('w.requires_grad:', w.requires_grad) # True\n",
    "print('o.requires_grad:', output.requires_grad) # True\n",
    "\n",
    "# their grad is not avaiable\n",
    "print('x.grad:', x.grad) # tensor([1., 2., 3., 4.])\n",
    "print('y.grad:', y.grad) # tensor([1., 2., 3., 4.])\n",
    "print('w.grad:', w.grad) # tensor([ 4.,  6.,  8., 10.])\n",
    "print('z.grad:', z.grad) # None\n",
    "print('o.grad:', output.grad) # None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. hook_fn(grad) -> Tensor or None\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Start backprop=====\n",
      "tensor([2., 4., 6., 8.])\n",
      "=====End backprop=====\n",
      "x.grad: tensor([2., 4., 6., 8.])\n",
      "y.grad: tensor([2., 4., 6., 8.])\n",
      "w.grad: tensor([ 4.,  6.,  8., 10.])\n",
      "z.grad: None\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:29: UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the gradient for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more informations.\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "x = torch.Tensor([0, 1, 2, 3]).requires_grad_()\n",
    "y = torch.Tensor([4, 5, 6, 7]).requires_grad_()\n",
    "w = torch.Tensor([1, 2, 3, 4]).requires_grad_()\n",
    "z = x+y\n",
    "\n",
    "# ===================\n",
    "def hook_fn(grad):\n",
    "    print(grad)\n",
    "\n",
    "def hook_fn_update(grad):\n",
    "    grad = 2 * grad\n",
    "    print(grad)\n",
    "    return grad\n",
    "\n",
    "z.register_hook(hook_fn_update)\n",
    "# ===================\n",
    "\n",
    "o = w.matmul(z)\n",
    "\n",
    "print('=====Start backprop=====')\n",
    "o.backward()\n",
    "print('=====End backprop=====')\n",
    "\n",
    "print('x.grad:', x.grad)\n",
    "print('y.grad:', y.grad)\n",
    "print('w.grad:', w.grad)\n",
    "print('z.grad:', z.grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Hook for Modules\n",
    "\n",
    "网络模块 module 不像上一节中的 Tensor，拥有显式的变量名可以直接访问，而是被封装在神经网络中间。我们通常只能获得网络整体的输入和输出，对于夹在网络中间的模块，我们不但很难得知它输入/输出的梯度,为了解决这个麻烦，PyTorch 设计了两种 hook.\n",
    "\n",
    "+ ``register_forward_hook``: **获取/修改**前向传播过程中，各个网络模块的输入和输出，**不返回任何值**\n",
    "+ ``register_backward_hook``： 获取神经网络反向传播过程中，各个模块***输入端***和***输出端的梯度值***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Linear(in_features=3, out_features=4, bias=True)\n",
      "input (tensor([[1., 1., 1.]], requires_grad=True),)\n",
      "output tensor([[  7., -13.,  27., -29.]], grad_fn=<AddmmBackward>)\n",
      "ReLU()\n",
      "input (tensor([[  7., -13.,  27., -29.]], grad_fn=<AddmmBackward>),)\n",
      "output tensor([[ 7.,  0., 27.,  0.]], grad_fn=<ReluBackward0>)\n",
      "Linear(in_features=4, out_features=1, bias=True)\n",
      "input (tensor([[ 7.,  0., 27.,  0.]], grad_fn=<ReluBackward0>),)\n",
      "output tensor([[89.]], grad_fn=<AddmmBackward>)\n",
      "Linear(in_features=4, out_features=1, bias=True)\n",
      "grad input  (tensor([1.]), tensor([[1., 2., 3., 4.]]), tensor([[ 7.],\n",
      "        [ 0.],\n",
      "        [27.],\n",
      "        [ 0.]]))\n",
      "grad ouput  (tensor([[1.]]),)\n",
      "ReLU()\n",
      "grad input  (tensor([[1., 0., 3., 0.]]),)\n",
      "grad ouput  (tensor([[1., 2., 3., 4.]]),)\n",
      "Linear(in_features=3, out_features=4, bias=True)\n",
      "grad input  (tensor([1., 0., 3., 0.]), tensor([[22., 26., 30.]]), tensor([[1., 0., 3., 0.],\n",
      "        [1., 0., 3., 0.],\n",
      "        [1., 0., 3., 0.]]))\n",
      "grad ouput  (tensor([[1., 0., 3., 0.]]),)\n",
      "\n",
      "\n",
      "==========Saved inputs and outputs==========\n",
      "input:  (tensor([[1., 1., 1.]], requires_grad=True),)\n",
      "output:  tensor([[  7., -13.,  27., -29.]], grad_fn=<AddmmBackward>)\n",
      "input:  (tensor([[  7., -13.,  27., -29.]], grad_fn=<AddmmBackward>),)\n",
      "output:  tensor([[ 7.,  0., 27.,  0.]], grad_fn=<ReluBackward0>)\n",
      "input:  (tensor([[ 7.,  0., 27.,  0.]], grad_fn=<ReluBackward0>),)\n",
      "output:  tensor([[89.]], grad_fn=<AddmmBackward>)\n",
      "\n",
      "\n",
      "==========Saved backward grad==========\n",
      "input:  (tensor([1.]), tensor([[1., 2., 3., 4.]]), tensor([[ 7.],\n",
      "        [ 0.],\n",
      "        [27.],\n",
      "        [ 0.]]))\n",
      "output:  (tensor([[1.]]),)\n",
      "input:  (tensor([[1., 0., 3., 0.]]),)\n",
      "output:  (tensor([[1., 2., 3., 4.]]),)\n",
      "input:  (tensor([1., 0., 3., 0.]), tensor([[22., 26., 30.]]), tensor([[1., 0., 3., 0.],\n",
      "        [1., 0., 3., 0.],\n",
      "        [1., 0., 3., 0.]]))\n",
      "output:  (tensor([[1., 0., 3., 0.]]),)\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "\n",
    "class Model(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Model, self).__init__()\n",
    "        self.fc1 = nn.Linear(3, 4)\n",
    "        self.relu1 = nn.ReLU()\n",
    "        self.fc2 = nn.Linear(4, 1)\n",
    "        self.initialize()\n",
    "    \n",
    "    def initialize(self):\n",
    "        with torch.no_grad():\n",
    "            self.fc1.weight = torch.nn.Parameter(\n",
    "                torch.Tensor([[1., 2., 3.],\n",
    "                              [-4., -5., -6.],\n",
    "                              [7., 8., 9.],\n",
    "                              [-10., -11., -12.]]))\n",
    "\n",
    "            self.fc1.bias = torch.nn.Parameter(torch.Tensor([1.0, 2.0, 3.0, 4.0]))\n",
    "            self.fc2.weight = torch.nn.Parameter(torch.Tensor([[1.0, 2.0, 3.0, 4.0]]))\n",
    "            self.fc2.bias = torch.nn.Parameter(torch.Tensor([1.0]))\n",
    "    \n",
    "    def forward(self, x):\n",
    "        o = self.fc1(x)\n",
    "        o = self.relu1(o)\n",
    "        o = self.fc2(o)\n",
    "        return o\n",
    "\n",
    "# 全局变量，存储中间 features\n",
    "total_feat_out = []\n",
    "total_feat_in = []\n",
    "\n",
    "def hook_fn_forward(module, input, output):\n",
    "    print(module) # 用于区分模块\n",
    "    print('input', input) # 首先打印出来\n",
    "    print('output', output)\n",
    "    total_feat_out.append(output) # 然后分别存入全局 list 中\n",
    "    total_feat_in.append(input)\n",
    "\n",
    "    \n",
    "total_grad_out = []\n",
    "total_grad_in = []\n",
    "def hook_fn_backward(module, grad_input, grad_output):\n",
    "    print(module)\n",
    "    print('grad input ', grad_input)\n",
    "    print('grad ouput ', grad_output)\n",
    "    total_grad_in.append(grad_input)\n",
    "    total_grad_out.append(grad_output)\n",
    "\n",
    "model = Model()\n",
    "for name, module in model.named_children():\n",
    "    module.register_forward_hook(hook_fn_forward)\n",
    "    module.register_backward_hook(hook_fn_backward)\n",
    "\n",
    "x = torch.Tensor([[1.0, 1.0, 1.0]]).requires_grad_() \n",
    "o = model(x)\n",
    "o.backward()\n",
    "\n",
    "print('\\n\\n==========Saved inputs and outputs==========')\n",
    "for idx in range(len(total_feat_in)):\n",
    "    print('input: ', total_feat_in[idx])\n",
    "    print('output: ', total_feat_out[idx])\n",
    "    \n",
    "print('\\n\\n==========Saved backward grad==========')\n",
    "for idx in range(len(total_feat_in)):\n",
    "    print('input: ', total_grad_in[idx])\n",
    "    print('output: ', total_grad_out[idx])"
   ]
  }
 ],
 "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": 4
}