Hook Management - part 2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hook Management - part 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### XLM - Multi30k machine translation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import transformers\n",
    "from transformers import XLMTokenizer, XLMWithLMHeadModel\n",
    "import spacy\n",
    "import torchtext\n",
    "from torchtext.datasets import Multi30k\n",
    "from torchtext.data import Field, BucketIterator\n",
    "import torch.nn as nn\n",
    "from tacklebox.hook_management import HookManager"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Python 3.7.3\n",
      "torch 1.2.0\n",
      "spacy 2.1.9\n",
      "transformers 2.4.1\n",
      "torchtext 0.4.0\n"
     ]
    }
   ],
   "source": [
    "!python --version\n",
    "print('torch %s' % torch.__version__)\n",
    "print('spacy %s' % spacy.__version__)\n",
    "print('transformers %s' % transformers.__version__)\n",
    "print('torchtext %s' % torchtext.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "english = spacy.load('en')\n",
    "german = spacy.load('de')\n",
    "\n",
    "def tokenize_en(text):\n",
    "    return [tok.text for tok in english.tokenizer(text)]\n",
    "def tokenize_de(text):\n",
    "    return [tok.text for tok in german.tokenizer(text)]\n",
    "\n",
    "en_text = Field(sequential=True, use_vocab=True, tokenize=tokenize_en, lower=True)\n",
    "de_text = Field(sequential=True, use_vocab=True, tokenize=tokenize_de, lower=True)\n",
    "\n",
    "train, val, test = Multi30k.splits(root='../data', exts=('.en', '.de'), fields=(en_text, de_text))\n",
    "\n",
    "en_text.build_vocab(train, max_size=30000, min_freq=3)\n",
    "de_text.build_vocab(train, max_size=30000, min_freq=3)\n",
    "vocab_en = en_text.vocab\n",
    "vocab_de = de_text.vocab\n",
    "pad_idx = vocab_de.stoi['<pad>']\n",
    "\n",
    "train_ldr, val_ldr, test_ldr = BucketIterator.splits((train, val, test),\n",
    "                                                    batch_size=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "xlm = XLMWithLMHeadModel.from_pretrained('xlm-mlm-ende-1024')\n",
    "xlm.transformer.embeddings = nn.Embedding(len(vocab_en), xlm.config.emb_dim, padding_idx=pad_idx)\n",
    "xlm.pred_layer.proj = nn.Linear(xlm.config.emb_dim, len(vocab_de), bias=True)\n",
    "_ = xlm.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "xent = nn.CrossEntropyLoss()\n",
    "\n",
    "batch = next(iter(train_ldr))\n",
    "src, trg = batch.src.to(0), batch.trg.to(0)\n",
    "\n",
    "def mt_loss(out, target):\n",
    "    # only compute loss for non-padding indices\n",
    "    min_idx = min([out.shape[0], target.shape[0]])\n",
    "    out, target = out[:min_idx], target[:min_idx]\n",
    "    mask = (target != pad_idx).type(torch.bool)\n",
    "    return xent(out[mask], target[mask])\n",
    "\n",
    "out, = xlm(src)\n",
    "mt_loss(out, trg).backward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import and initialize the hook manager\n",
    "hookmngr = HookManager()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define named modules, including the embedding layer (id=embedding) and final attention (id=final_attn)\n",
    "named_modules = {\n",
    "    'embedding': xlm.transformer.embeddings,\n",
    "    'final_attn': xlm.transformer.attentions[-1]\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Filtering hooks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# forward pre-hook function signature: (module, inputs)\n",
    "def zero_input(module, inputs):\n",
    "    ret = []\n",
    "    zeroed = 0\n",
    "    for input in inputs:\n",
    "        if type(input) == torch.Tensor and input.dtype == torch.float:\n",
    "            ret += [input - input]\n",
    "            zeroed += 1\n",
    "        else:\n",
    "            ret += [input]\n",
    "    print('Set %d/%d of %s inputs to zero' % (zeroed, len(inputs), module.name))\n",
    "    return tuple(ret)\n",
    "\n",
    "def print_mean(module, inputs, outputs):\n",
    "    print('%s input mean = %.2f, output mean = %.2f' % (module.name,\n",
    "                                                        inputs[0].sum().item() / inputs[0].numel(),\n",
    "                                                        outputs[0].sum().item() / outputs[0].numel()))\n",
    "\n",
    "# register both hooks on all named modules, deferring activation\n",
    "hookmngr.register_forward_pre_hook(zero_input, hook_fn_name='zero_input', activate=False, **named_modules)\n",
    "hookmngr.register_forward_hook(print_mean, hook_fn_name='print_mean', activate=False, **named_modules)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Set 0/1 of embedding inputs to zero\n",
      "embedding input mean = 101.78, output mean = 0.00\n",
      "Set 1/2 of final_attn inputs to zero\n",
      "final_attn input mean = 0.00, output mean = -0.00\n"
     ]
    }
   ],
   "source": [
    "# activate all hooks\n",
    "with hookmngr.hook_all_context() + torch.no_grad():\n",
    "    xlm(src)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "embedding input mean = 101.78, output mean = 0.00\n",
      "final_attn input mean = -0.01, output mean = 0.02\n"
     ]
    }
   ],
   "source": [
    "# lets only activate our forward hooks\n",
    "with hookmngr.hook_all_context(category='forward_hook') + torch.no_grad():\n",
    "    xlm(src)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "embedding input mean = 101.78, output mean = 0.00\n",
      "final_attn input mean = -0.01, output mean = 0.02\n"
     ]
    }
   ],
   "source": [
    "# now lets use only our original hook function, print_mean\n",
    "with hookmngr.hook_all_context(hook_types=[print_mean]) + torch.no_grad():\n",
    "    xlm(src)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Backward hooks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# backward hooks\n",
    "def print_grad_shape(module, grad_in, grad_out):\n",
    "    print('%s ' % module.name, end='')\n",
    "    if grad_in[0] is not None:\n",
    "        print('grad_in shape: ', grad_in[0].shape, end=', ')\n",
    "    if grad_out[0] is not None:\n",
    "        print('grad_out shape: ', grad_out[0].shape, end='')\n",
    "    print('')\n",
    "\n",
    "# register backward hook print_grad_shape on the final attentionlayer\n",
    "hookmngr.register_backward_hook(print_grad_shape, named_modules['final_attn'],\n",
    "                               activate=False, hook_fn_name='print_grad_shape')\n",
    "\n",
    "# Note: can pass modules as additional args if already registered"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "final_attn grad_in shape:  torch.Size([19, 5, 1024]), grad_out shape:  torch.Size([19, 5, 1024])\n"
     ]
    }
   ],
   "source": [
    "# activate all backward hooks\n",
    "with hookmngr.hook_all_context(category='backward_hook'):\n",
    "    out, = xlm(src)\n",
    "    mt_loss(out, trg).backward()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using inputs, outputs and gradients all at once"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# using intermediate outputs with gradient\n",
    "def print_outputs_with_grad(module, grad_in, grad_out, inputs, outputs):\n",
    "    print('%s input-gradient pairs: ' % module.name, end='')\n",
    "    for inp, grad in zip(inputs, grad_in):\n",
    "        print(inp.dtype, type(grad), end=', ')\n",
    "    print('')\n",
    "    print('%s output-gradient pairs: ' % module.name, end='')\n",
    "    for out, grad in zip(outputs, grad_out):\n",
    "        print(out.dtype, type(grad), end=', ')\n",
    "    print('')\n",
    "\n",
    "# use retain_forward_cache argument to provide hook function access to forward pass data during backward pass\n",
    "hookmngr.register_backward_hook(print_outputs_with_grad, named_modules['final_attn'],\n",
    "                               activate=False, hook_fn_name='print_outputs_with_grad',\n",
    "                               retain_forward_cache=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "final_attn input-gradient pairs: torch.float32 <class 'torch.Tensor'>, torch.bool <class 'NoneType'>, \n",
      "final_attn output-gradient pairs: torch.float32 <class 'torch.Tensor'>, \n"
     ]
    }
   ],
   "source": [
    "# activate only print_outputs_with_grad\n",
    "with hookmngr.hook_all_context(hook_types=[print_outputs_with_grad]):\n",
    "    out, = xlm(src)\n",
    "    mt_loss(out, trg).backward()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hook removal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "zero_input[embedding]\n",
      "zero_input[final_attn]\n",
      "print_mean[embedding]\n",
      "print_mean[final_attn]\n",
      "print_grad_shape[final_attn]\n",
      "print_outputs_with_grad[final_attn]\n"
     ]
    }
   ],
   "source": [
    "print('\\n'.join(hookmngr.name_to_hookhandle.keys()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# remove print_mean from final_attn\n",
    "hookmngr.remove_hook_by_name('print_mean[final_attn]')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "zero_input[embedding]\n",
      "zero_input[final_attn]\n",
      "print_mean[embedding]\n",
      "print_grad_shape[final_attn]\n",
      "print_outputs_with_grad[final_attn]\n"
     ]
    }
   ],
   "source": [
    "print('\\n'.join(hookmngr.name_to_hookhandle.keys()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# remove zero_input from all modules\n",
    "hookmngr.remove_hook_function(zero_input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "print_mean[embedding]\n",
      "print_grad_shape[final_attn]\n",
      "print_outputs_with_grad[final_attn]\n"
     ]
    }
   ],
   "source": [
    "print('\\n'.join(hookmngr.name_to_hookhandle.keys()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# remove all remaining hooks from final_attn\n",
    "hookmngr.remove_module_by_name('final_attn')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "print_mean[embedding]\n"
     ]
    }
   ],
   "source": [
    "print('\\n'.join(hookmngr.name_to_hookhandle.keys()))"
   ]
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}