attention_exercise.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Network attention\n",
    "\n",
    "In this tutorial we will learn how to analyse the attention of a network model. We are going to: \n",
    "\n",
    "1) Derive saliency maps\n",
    "\n",
    "2) Compare saliency maps for different inputs\n",
    "\n",
    "3) (optional) Derive class activation maps and compare\n",
    "\n",
    "### Data and trained network model\n",
    "The data are simulated astrophysical 3D lightcones of the 21cm signal of the forbidden spin-flip transition of neutral hydrogen, generated with the semi-numerical code 21cmFAST (https://github.com/21cmfast/21cmFAST), in the case of mock lightcones noise was generated with the code 21cmSense (https://github.com/jpober/21cmSense), to represent data as expected for the Square Kilometre Array at radio wavelengths. The trained network model is the 3D-21cmPIE-Net (see https://github.com/stef-neu/3D-21cmPIE-Net) as published in Neutsch, Heneka, Brueggen 2022 (arXiv:2201.07587). Credit of the code go to Steffen Neutsch, Caroline Heneka 2022. \n",
    "\n",
    "### General remarks\n",
    "As we use trained models for this task, GPU are not necessarily required."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# some packages that we will need\n",
    "#uncomment on google colab:\n",
    "# !git clone https://github.com/csheneka/introspection-tutorial.git\n",
    "#!pip install tensorflow==2.9.1\n",
    "#!pip install tf-keras-vis tensorflow\n",
    "\n",
    "import tensorflow as tf\n",
    "print(tf.__version__)\n",
    "from tensorflow import keras\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import os\n",
    "from matplotlib import cm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# optional color-scheme for plotting later on\n",
    "from matplotlib import colors\n",
    "eor_colour = colors.LinearSegmentedColormap.from_list(\n",
    "    \"EoR\",\n",
    "    [\n",
    "        (0, \"white\"),\n",
    "        (0.21, \"yellow\"),\n",
    "        (0.42, \"orange\"),\n",
    "        (0.63, \"red\"),\n",
    "        (0.86, \"black\"),\n",
    "        (0.9, \"blue\"),\n",
    "        (1, \"cyan\"),\n",
    "    ],\n",
    ")\n",
    "plt.register_cmap(cmap=eor_colour)\n",
    "\n",
    "# Change the sigmoid activation of the last layer to a linear one to not obstruct attention.\n",
    "def model_modifier(m):\n",
    "    m.layers[-1].activation = tf.keras.activations.linear\n",
    "    return m"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1) Derive saliency maps\n",
    "\n",
    "To derive saliency, we need to load both model and data. The trained model is available both trained on simulations-only and on mock data including noise. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the trained model(s)\n",
    "# model 1 simulations-only: models/3D_21cmPIE_Net_sim_par6.h5\n",
    "# model 2 mock: models/3D_21cmPIE_Net_optmock_par6.h5\n",
    "model = keras.models.load_model('models/3D_21cmPIE_Net_sim_par6.h5')\n",
    "#google colab: one might have to adjust the path e.g. to 'introspection-tutorial/models/3D_21cmPIE_Net_sim_par6.h5'\n",
    "\n",
    "# Read light-cones from tfrecords files\n",
    "HII_DIM=140\n",
    "BOX_LEN=2350\n",
    "nlabel = 6\n",
    "nz = 92\n",
    "def parse_function(files):\n",
    "    keys_to_features = {\"label\":tf.io.FixedLenFeature((nlabel),tf.float32),\n",
    "                        \"image\":tf.io.FixedLenFeature((),tf.string),\n",
    "                        \"tau\":tf.io.FixedLenFeature((),tf.float32),\n",
    "                        \"gxH\":tf.io.FixedLenFeature((nz),tf.float32),\n",
    "                        \"z\":tf.io.FixedLenFeature((nz),tf.float32),}\n",
    "    parsed_features = tf.io.parse_example(files, keys_to_features)\n",
    "    image = tf.io.decode_raw(parsed_features[\"image\"],tf.float16)\n",
    "    image = tf.reshape(image,(HII_DIM,HII_DIM,BOX_LEN))\n",
    "    return image, parsed_features[\"label\"] # Image, label ordering: m_WDM,Omega_m,L_X,E_0,T_vir,zeta\n",
    "\n",
    "data_file=\"data/lc0.tfrecord\"  # larger file: lc1n.tfrecord\n",
    "#google colab: one might have to adjust the path e.g. to 'introspection-tutorial/data/lc0.tfrecord'\n",
    "dataset = tf.data.TFRecordDataset(data_file)\n",
    "dataset = dataset.map(parse_function,num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
    "data = np.array(list(dataset.as_numpy_iterator()),dtype=object)\n",
    "print(data.shape)\n",
    "print(data[0,0].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now move on to derive saliency, using the tf-keras-vis package.\n",
    "\n",
    "### Task 1: \n",
    "Explore plotting saliency maps, how do they look, what do they tell?\n",
    "\n",
    "### Question 1: \n",
    "What do you notice, which areas are spatially and temporally important?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tf_keras_vis.saliency import Saliency\n",
    "\n",
    "# create saliency object\n",
    "obj_saliency = Saliency(model, model_modifier=model_modifier,clone=True) \n",
    "\n",
    "# generate saliency maps for parameters i.e. classes\n",
    "parameters=[0,1,2,3,4,5]\n",
    "lc = data[0,0]\n",
    "for para in parameters:\n",
    "    def loss(output):\n",
    "        return output[0][para] # shape (samples,classes)\n",
    "    map_saliency = obj_saliency(loss,lc.reshape(?,?,?,1))\n",
    "\n",
    "print(map_saliency.shape)\n",
    "nslice = \n",
    "plt.figure()\n",
    "plot_saliency = plt.imshow(map_saliency[?], cmap=cm.hot)\n",
    "plt.figure()\n",
    "plot_lc = plt.imshow(lc[nslice,:,:], cmap=\"EoR\",vmin=-150,vmax=30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2a) Compare saliency maps for different inputs\n",
    "\n",
    "### Task 2a: \n",
    "Plot both saliency for simulations-only and mock lightcones\n",
    "\n",
    "### Question 2a: \n",
    "Does the network's attention shift with the inclusion of noise? If yes, how?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Task 2, plot and compare here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2b) Beyond vanilla saliency \n",
    "### Question 2b: \n",
    "Can the mapping of attention with saliency be improved, in particular can noise in the maps be reduced? Tipp: Use SmoothGrad from tf-keras-vis. What does SmoothGrad exploit?\n",
    "\n",
    "### Task 2b:\n",
    "Explore SmoothGrad for different setting (smooth_samples, smooth_noise)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#  Question 2b, try out SmoothGrad here\n",
    "from tf_keras_vis.saliency import Saliency\n",
    "\n",
    "# Change the sigmoid activation of the last layer to a linear one to not obstruct attention.\n",
    "def model_modifier(m):\n",
    "    m.layers[-1].activation = tf.keras.activations.linear\n",
    "    return m\n",
    "\n",
    "# create saliency object\n",
    "obj_saliency = Saliency(?) \n",
    "\n",
    "nsamples = \n",
    "noise = \n",
    "\n",
    "# generate saliency maps for parameters i.e. classes\n",
    "parameters=[0,1,2,3,4,5]\n",
    "lc = data[0,0]\n",
    "for para in parameters:\n",
    "    def loss(output):\n",
    "        return output[0][para] # shape (samples,classes)\n",
    "    map_saliency = obj_saliency(?(see above),?(see above),smooth_samples=nsamples,smooth_noise=noise)\n",
    "\n",
    "print(map_saliency.shape)\n",
    "nslice =\n",
    "plt.figure()\n",
    "plot_saliency = plt.imshow(map_saliency[?], cmap=cm.hot)\n",
    "plt.figure()\n",
    "plot_lc = plt.imshow(lc[nslice,:,:], cmap=\"EoR\",vmin=-150,vmax=30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3) (optional) Derive class activation maps and compare \n",
    "\n",
    "### Task 3: \n",
    "Create CAMs, explore their spatial and temporal structure. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# note: \n",
    "# there is currently an incompatibility of CAM generated with tf-keras-vis and the newest tensorflow version 2.9\n",
    "# will need to downgrade tensorflow version\n",
    "from tf_keras_vis.gradcam import Gradcam\n",
    "\n",
    "# create CAM object\n",
    "obj_cam = gradcam.Gradcam(model, model_modifier=model_modifier,clone=True) \n",
    "\n",
    "# generate heatmaps for attention\n",
    "parameters=[0,1,2,3,4,5]\n",
    "lc = data[0,0]\n",
    "for para in parameters:\n",
    "    def loss(output):\n",
    "        return output[0][para] # shape (samples,classes)\n",
    "    map_cam = obj_cam(loss,lc.reshape(140,140,2350,1),penultimate_layer=-1) # focus on last convolutional layer\n",
    "\n",
    "print(map_cam.shape)\n",
    "nslice = 130\n",
    "plt.figure()\n",
    "plot_cam = plt.imshow(map_cam[0,nslice,:,:], cmap=cm.hot)\n",
    "plt.figure()\n",
    "plot_lc = plt.imshow(lc[nslice,:,:], cmap=\"EoR\",vmin=-150,vmax=30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}