source_target_contributions/1_Load_model_and_evaluate_LRP.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluate LRP for a trained model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "sys.path.insert(0, 'path_to_the-story-of-heads') # insert your local path to the repo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load vocabularies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import pickle\n",
    "import numpy as np\n",
    "\n",
    "VOC_PATH =  # insert your path\n",
    "\n",
    "inp_voc = pickle.load(open(VOC_PATH + 'src.voc', 'rb'))\n",
    "out_voc = pickle.load(open(VOC_PATH + 'dst.voc', 'rb'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env: CUDA_VISIBLE_DEVICES=0\n"
     ]
    }
   ],
   "source": [
    "%env CUDA_VISIBLE_DEVICES=0\n",
    "\n",
    "import tensorflow as tf\n",
    "import lib\n",
    "import lib.task.seq2seq.models.transformer_lrp as tr\n",
    "\n",
    "tf.reset_default_graph()\n",
    "gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.99, allow_growth=True)\n",
    "sess = tf.InteractiveSession(config=tf.ConfigProto(gpu_options=gpu_options))\n",
    "\n",
    "hp = {\n",
    "     \"num_layers\": 6,\n",
    "     \"num_heads\": 8,\n",
    "     \"ff_size\": 2048,\n",
    "     \"ffn_type\": \"conv_relu\",\n",
    "     \"hid_size\": 512,\n",
    "     \"emb_size\": 512,\n",
    "     \"res_steps\": \"nlda\", \n",
    "    \n",
    "     \"rescale_emb\": True,\n",
    "     \"inp_emb_bias\": True,\n",
    "     \"normalize_out\": True,\n",
    "     \"share_emb\": False,\n",
    "     \"replace\": 0,\n",
    "    \n",
    "     \"relu_dropout\": 0.1,\n",
    "     \"res_dropout\": 0.1,\n",
    "     \"attn_dropout\": 0.1,\n",
    "     \"label_smoothing\": 0.1,\n",
    "    \n",
    "     \"translator\": \"ingraph\",\n",
    "     \"beam_size\": 4,\n",
    "     \"beam_spread\": 3,\n",
    "     \"len_alpha\": 0.6,\n",
    "     \"attn_beta\": 0,\n",
    "}\n",
    "\n",
    "model = tr.Model('mod', inp_voc, out_voc, inference_mode='fast', **hp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "path_to_ckpt = # insert path to the final checkpoint\n",
    "var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)\n",
    "lib.train.saveload.load(path_to_ckpt, var_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Sanity check: translate something and see if the model is ok"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /opt/lena-voita/test_heads_repo/the-story-of-heads/lib/layers/basic.py:144: calling reduce_mean (from tensorflow.python.ops.math_ops) with keep_dims is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "keep_dims is deprecated, use keepdims instead\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[\"j' ai vu un cat .\"]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.translate_lines(['i saw a cat .'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note! Your dataset has to have the same length of all source sentences, as well as all target sentences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "datadir = # insert your path\n",
    "\n",
    "test_src = open(datadir + 'YOUR FILE NAME').readlines()\n",
    "test_dst = open(datadir + 'YOUR FILE NAME').readlines()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "feed_dict = model.make_feed_dict(zip(test_src[:3], test_dst[:3]))\n",
    "ph = lib.task.seq2seq.data.make_batch_placeholder(feed_dict)\n",
    "feed = {ph[key]: feed_dict[key] for key in feed_dict}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Functions to evaluate LRP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lib.ops.record_activations import recording_activations\n",
    "from lib.layers.basic import dropout_scope\n",
    "from lib.ops import record_activations as rec\n",
    "from lib.layers.lrp import LRP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_topk_logits_selector(logits, k=3):\n",
    "    \"\"\" takes logits[batch, nout, voc_size] and returns a mask with ones at k largest logits \"\"\"\n",
    "    topk_logit_indices = tf.nn.top_k(logits, k=k).indices\n",
    "    indices = tf.stack([\n",
    "        tf.range(tf.shape(logits)[0] * tf.shape(logits)[1] * k) // (tf.shape(logits)[1] * k),\n",
    "        (tf.range(tf.shape(logits)[0] * tf.shape(logits)[1] * k) // k) % tf.shape(logits)[1],\n",
    "        tf.reshape(topk_logit_indices, [-1])\n",
    "    ], axis=1)\n",
    "    ones = tf.ones(shape=(tf.shape(indices)[0],))\n",
    "    return tf.scatter_nd(indices, ones, shape=tf.shape(logits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /opt/lena-voita/test_heads_repo/the-story-of-heads/lib/layers/lrp.py:87: calling reduce_sum (from tensorflow.python.ops.math_ops) with keep_dims is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "keep_dims is deprecated, use keepdims instead\n"
     ]
    }
   ],
   "source": [
    "target_position = tf.placeholder(tf.int32, [])\n",
    "with rec.recording_activations() as saved_activations, dropout_scope(False):\n",
    "\n",
    "    rdo = model.encode_decode(ph, is_train=False)\n",
    "    logits = model.loss._rdo_to_logits(rdo)\n",
    "    out_mask = tf.one_hot(target_position, depth=tf.shape(logits)[1])[None, :, None]\n",
    "    \n",
    "    top1_logit = get_topk_logits_selector(logits, k=1) * tf.nn.softmax(logits)\n",
    "    top1_prob = tf.reduce_sum(top1_logit, axis=-1)[0]\n",
    "\n",
    "    R_ = get_topk_logits_selector(logits, k=1) * out_mask\n",
    "    R = model.loss._rdo_to_logits.relprop(R_)\n",
    "    R = model.transformer.relprop_decode(R)\n",
    "    \n",
    "    R_out = tf.reduce_sum(abs(R['emb_out']), axis=-1)\n",
    "    R_inp = tf.reduce_sum(abs(model.transformer.relprop_encode(R['enc_out'])), axis=-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate LRP for your dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "dir_out = '.' # set the directory to save the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = []\n",
    "\n",
    "for elem in zip(test_src, test_dst):\n",
    "    #print(len(result))\n",
    "    src = elem[0].strip()\n",
    "    dst = elem[1].strip()\n",
    "    dst_words = len(dst.split()) + 1\n",
    "    feed_dict = model.make_feed_dict(zip([src], [dst]))\n",
    "    feed = {ph[key]: feed_dict[key] for key in feed_dict}\n",
    "    \n",
    "    inp_lrp = []\n",
    "    out_lrp = []\n",
    "    for token_pos in range(feed_dict['out'].shape[1]):\n",
    "        feed[target_position] = token_pos\n",
    "        res_inp, res_out = sess.run((R_inp, R_out), feed)\n",
    "        inp_lrp.append(res_inp[0])\n",
    "        out_lrp.append(res_out[0])\n",
    "    result.append({'src': src, 'dst': dst,\n",
    "                   'inp_lrp': np.array(inp_lrp), 'out_lrp': np.array(out_lrp)\n",
    "                  })\n",
    "    \n",
    "pickle.dump(result, open(dir_out + 'lrp_results', 'wb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}