{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-36JeZ6mXluU"
   },
   "source": [
    "# Tensorboard Introduction\n",
    "\n",
    "Welcome to the introduction of [`TensorBoard`](https://www.tensorflow.org/tensorboard). In this tutorial, we’ll learn how to:\n",
    "\n",
    "1. Set up TensorBoard\n",
    "2. Write values to TensorBoard\n",
    "3. Inspect a model architecture using TensorBoard\n",
    "4. Train model and write loss, accuracy and some images to TensorBoard\n",
    "\n",
    "Finally we will visualize the effect of different weight initializations on the neural network using `TensorBoard`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "x7p2c7L2XluW"
   },
   "source": [
    "## (Optional) Mount folder in Colab\n",
    "\n",
    "Uncomment the following cell to mount your gdrive if you are using the notebook in google colab:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-04-16T10:59:42.597478Z",
     "start_time": "2023-04-16T10:59:42.146691Z"
    },
    "id": "KeYPBZngXluY"
   },
   "outputs": [],
   "source": [
    "# Optional: install correct libraries in google colab. Again, downgrade and then RESTART the session! Runtime > Restart Session.\n",
    "# !python -m pip install \"numpy<2.0\" torch==2.2.2 torchvision==0.17.2 torchaudio==2.2.2 --index-url https://download.pytorch.org/whl/cu121\n",
    "# !python -m pip install numpy==1.26.4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-04-16T10:59:42.572717Z",
     "start_time": "2023-04-16T10:59:42.120637Z"
    },
    "id": "oYkKZpz1XluX"
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "from google.colab import drive\n",
    "import os\n",
    "\n",
    "gdrive_path='/content/gdrive/MyDrive/i2dl/exercise_07'\n",
    "\n",
    "# This will mount your google drive under 'MyDrive'\n",
    "drive.mount('/content/gdrive', force_remount=True)\n",
    "# In order to access the files in this notebook we have to navigate to the correct folder\n",
    "os.chdir(gdrive_path)\n",
    "# Check manually if all files are present\n",
    "print(sorted(os.listdir()))\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vRELjWdMXluZ"
   },
   "source": [
    "# 1. Getting Started"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XoYT1xqVXlua"
   },
   "source": [
    "TensorBoard helps us track our metrics such as loss, accuracy and visualize the results, model graphs that may be needed during the machine learning workflow.\n",
    "\n",
    "Let's start by installing `TensorBoard` and maybe the correct alternative libraries for google colab.\n",
    "\n",
    "# **Note**:\n",
    "All packages should be installed on your i2dl conda enviroment. Otherwise, you start with mismatching versions loops of differnet libraries, which will make your life really difficult later on.\n",
    "\n",
    "There might be a warning regarding TensorFlow compatibility, but it can be safely ignored. The setup should still work correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-04-16T10:59:42.599334Z",
     "start_time": "2023-04-16T10:59:42.170900Z"
    },
    "id": "O2VnH8UiXlua"
   },
   "outputs": [],
   "source": [
    "# remove the \"> /dev/null\" if you want to see the installation output\n",
    "import sys\n",
    "\n",
    "# For google colab\n",
    "# !python -m pip install tensorboard==2.9.1 > /dev/null\n",
    "\n",
    "# For anaconda/regular environments\n",
    "!{sys.executable} -m pip install tensorboard==2.9.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dcZ8LiXbXlua"
   },
   "outputs": [],
   "source": [
    "from exercise_code.tests.base_tests import bcolors\n",
    "import tensorboard\n",
    "import torch\n",
    "import torchvision\n",
    "print(f\"Tensorboard version: {tensorboard.__version__}\")\n",
    "if not tensorboard.__version__.startswith(\"2.9\"):\n",
    "    print(bcolors.colorize(\"red\", \"WARNING:\"),\n",
    "          bcolors.colorize(\"yellow\", \"You are using an another version of Tensorboard. We expect Tensorboard 2.9.1 You may continue using your version but it\"\n",
    "          \" might cause dependency and compatibility issues.\"))\n",
    "else:\n",
    "    print(bcolors.colorize(\"green\", \"SUCCESS:\"),\n",
    "          bcolors.colorize(\"yellow\", f\"You are using the correct version of Tensorboard: {bcolors.colorize('blue', tensorboard.__version__)}\"))\n",
    "\n",
    "\n",
    "if not torch.__version__.startswith(\"2.2\"):\n",
    "    print(bcolors.colorize(\"red\", \"WARNING:\"),\n",
    "          bcolors.colorize(\"yellow\", \"You are using an another version of PyTorch. We expect PyTorch 2.2.2. You may continue using your version but it \\\n",
    "              might cause dependency and compatibility issues.\"))\n",
    "else:\n",
    "    print(bcolors.colorize(\"green\", \"SUCCESS:\"),\n",
    "          bcolors.colorize(\"yellow\", f\"You are using the correct version of PyTorch: {bcolors.colorize('blue', torch.__version__)}\"))\n",
    "\n",
    "\n",
    "if not torchvision.__version__.startswith(\"0.17\"):\n",
    "    print(bcolors.colorize(\"red\", \"WARNING:\"),\n",
    "          bcolors.colorize(\"yellow\", \"you are using an another version of torchvision. We expect torchvision 0.17.2. You can continue with your version but it \\\n",
    "              might cause dependency and compatibility issues.\"))\n",
    "else:\n",
    "    print(bcolors.colorize(\"green\", \"SUCCESS:\"),\n",
    "          bcolors.colorize(\"yellow\", f\"You are using the correct version of torchvision: {bcolors.colorize('blue', torchvision.__version__)}\"))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IkzvlQVSXlub"
   },
   "source": [
    "This tutorial is highly aligned with the `TensorBoard` tutorial from [`PyTorch`](https://pytorch.org/tutorials/intermediate/tensorboard_tutorial.html). Have a look at that tutorial as well!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OPl9FHUmXluc"
   },
   "source": [
    "## 2. Setting up  TensorBoard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0_AOKYGnXlud"
   },
   "source": [
    "Let's start from where we ended the previous notebook on `PyTorch`. We will again use the [`Fashion-MNIST`](https://github.com/zalandoresearch/fashion-mnist) dataset for this notebook.\n",
    "\n",
    "The below cell of code sets up the dataloader and a plotting function to visualize samples from the dataset. This step is very similar to our previous notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qj_mwAhvXlud"
   },
   "outputs": [],
   "source": [
    "# import all the required packages\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "import os\n",
    "import shutil\n",
    "\n",
    "transform = transforms.Compose([transforms.ToTensor(),\n",
    "                                transforms.Normalize((0.5,),(0.5,))])  # mean and std have to be sequences (e.g. tuples),\n",
    "                                                                      # therefore we should add a comma after the values\n",
    "\n",
    "fashion_mnist_dataset = torchvision.datasets.FashionMNIST(root='../datasets', train=True,\n",
    "                                                          download=True, transform=transform)\n",
    "\n",
    "fashion_mnist_test_dataset = torchvision.datasets.FashionMNIST(root='../datasets', train=False,\n",
    "                                                          download=True, transform=transform)\n",
    "\n",
    "trainloader = torch.utils.data.DataLoader(fashion_mnist_dataset, batch_size=8)\n",
    "\n",
    "testloader = torch.utils.data.DataLoader(fashion_mnist_test_dataset, batch_size=8)\n",
    "\n",
    "classes = ('T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',\n",
    "           'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot')\n",
    "\n",
    "\n",
    "def matplotlib_imshow(img, one_channel=False):\n",
    "    if one_channel:\n",
    "        img = img.cpu().mean(dim=0)\n",
    "    img = img / 2 + 0.5     # unnormalize\n",
    "    npimg = img.numpy()\n",
    "    if one_channel:\n",
    "        plt.imshow(npimg, cmap=\"Greys\")\n",
    "    else:\n",
    "        plt.imshow(np.transpose(npimg, (1, 2, 0)))\n",
    "\n",
    "os.environ['KMP_DUPLICATE_LIB_OK']='True' # To prevent the kernel from dying.\n",
    "\n",
    "path = os.path.abspath(\"logs\")\n",
    "if os.path.exists(path):\n",
    "    shutil.rmtree(path)\n",
    "os.makedirs(path, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9KxntSZPXlud"
   },
   "source": [
    "Always remember to initialize the `device` variable with CUDA enabled GPU, in case it is available. This makes porting of our code to GPU's easier later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ompms8PCXlue"
   },
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(\"Using the device\",device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SuOGVVI0Xluf"
   },
   "source": [
    "Let us now intialize a 2-layer neural network model using the `nn.Module` of PyTorch. The model is then moved to the device specified by the `device` variable. We also complete the definitions of the loss function and the optimizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8OcVdzy1Xluf"
   },
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self, activation=nn.Sigmoid(),\n",
    "                 input_size=1*28*28, hidden_size=100, classes=10):\n",
    "        super().__init__()\n",
    "        self.input_size = input_size\n",
    "\n",
    "        # Here we initialize our activation and set up our two linear layers\n",
    "        self.activation = activation\n",
    "        self.fc1 = nn.Linear(input_size, hidden_size)\n",
    "        self.fc2 = nn.Linear(hidden_size, classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.view(-1, self.input_size) # flatten\n",
    "        x = self.fc1(x)\n",
    "        x = self.activation(x)\n",
    "        x = self.fc2(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "net = Net()\n",
    "net.to(device)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3Kx1SeC1Xluf"
   },
   "source": [
    "PyTorch provides support for logging data to TensorBoard using the `SummaryWriter` module.\n",
    "We will now initialize an object of`SummaryWriter` and specify the directory [**logs/introduction**] to store its related data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EGP3ZAJOXluf"
   },
   "outputs": [],
   "source": [
    "from torch.utils.tensorboard import SummaryWriter\n",
    "\n",
    "writer = SummaryWriter('logs/introduction')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bkcs0o_KXluf"
   },
   "source": [
    "# 3. Writing to TensorBoard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "D4jSBUSBXlug"
   },
   "source": [
    "Let's write some stuff to TensorBoard, and log into it to see how things go. :)\n",
    "\n",
    "\n",
    "1. You can open the TensorBoard GUI by running the command from this exercise folder in a Terminal:\n",
    "```tensorboard --logdir=./```\n",
    "\n",
    "For those using Linux or Mac, you can open a Terminal **in this exercise folder** and run the above command.\n",
    "\n",
    "For those using Windows with Anaconda packages, open an Anaconda Prompt and then run the above command.\n",
    "In case you don't  use Anaconda, use your default method of running python code in cmd.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "h7PhcgtJXlug"
   },
   "source": [
    "![tensorBoard Terminal](https://i2dl.vc.in.tum.de/static/images/exercise_07/tb_terminal.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KCuxK1oNXlug"
   },
   "source": [
    "You must be able to see the URL link (  `http://localhost:6006/` in the image) for accessing the tensorboard interface. Let's navigate to that URL in a browser.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mTd9l_cHXlug"
   },
   "source": [
    "2) Also, if you're using VScode, press <code> Ctrl + Shift + p (Windows User), CMD + Shift + P (Mac User)</code> and look for <code> Tensorboard </code> (Make sure that version <code>2.8.0</code> is installed)\n",
    "\n",
    "![](https://i2dl.vc.in.tum.de/static/images/exercise_07/Vscode_tensorboard.png)\n",
    "![](https://i2dl.vc.in.tum.de/static/images/exercise_07/Vscode_tensorboard_2.png)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "88zXdGjmXlug"
   },
   "outputs": [],
   "source": [
    "################## COLAB ##################\n",
    "# This might also work with jupyter notebooks, but will most likely not function well. Use the CMD/Terminal if possible (tensorboard --logdir=./)\n",
    "\n",
    "# %load_ext tensorboard\n",
    "\n",
    "# %tensorboard --logdir=./"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mMFCgaD2Xluh"
   },
   "source": [
    "\n",
    "No dashboards are created yet! Let's log some data to our `SummaryWriter` object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IJNGmJAWXluh"
   },
   "outputs": [],
   "source": [
    "# get some random training images\n",
    "for images, labels in trainloader:\n",
    "    break\n",
    "\n",
    "# create grid of images\n",
    "img_grid = torchvision.utils.make_grid(images)\n",
    "\n",
    "# show images using our helper function\n",
    "matplotlib_imshow(img_grid)\n",
    "\n",
    "# Write the generated image to tensorboard\n",
    "writer.add_image('four_mnist_images', img_grid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vbU9gga_Xluh"
   },
   "source": [
    "We can now see the image in our TensorBoard interface.  You might need to hit the refresh button on the top right as TensorBoard will only update in discrete intervals of time."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lZVT8PBiXluh"
   },
   "source": [
    "![tensorBoard Interface](./images/imgvis.jpg)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WuXDCgBmXluh"
   },
   "source": [
    "# 4. Visualization Model Architectures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6SU3nAwzXlum"
   },
   "source": [
    "Let's try to now visualize the architecture of our `net` model in Tensorboard. We can even look at input and output dimensions of your model. It is also a good way to debug as the model grows more and more complex.  \n",
    "\n",
    "Let's visualize the model now.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8BJ09yIaXlum"
   },
   "outputs": [],
   "source": [
    "writer.add_graph(net.cpu(), images)\n",
    "writer.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "30HiyOPNXlum"
   },
   "source": [
    "![Model Architecture Visualization](https://i2dl.vc.in.tum.de/static/images/exercise_07/tb_model.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9rTUvf_IXlun"
   },
   "source": [
    "Click the `GRAPHS` section in the top ribbon to access it the architecture. The above image was generated by clicking on our network `Net`.\n",
    "Feel free to explore with the various features of this model's visualization!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "57BTw3TfXlun"
   },
   "source": [
    "# 5. Training network models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aE6oBvwZXlun"
   },
   "source": [
    "It's now time to explore the most important use of TensorBoard - for model training.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vpAt6oBNXlun"
   },
   "source": [
    "We shall define two helper functions here: `images_to_probs` and `plot_classes_preds`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FpJnwqivXlun"
   },
   "outputs": [],
   "source": [
    "def images_to_probs(net, images):\n",
    "    '''\n",
    "    Returns the predicted class and probabilites of the image belonging to each of the classes\n",
    "    from the network output\n",
    "    '''\n",
    "    output = net(images)\n",
    "    # convert output probabilities to predicted class\n",
    "    _, preds_tensor = torch.max(output, 1)\n",
    "    preds = np.squeeze(preds_tensor.cpu().numpy())\n",
    "    return preds, [F.softmax(el, dim=0)[i].item() for i, el in zip(preds, output)]\n",
    "\n",
    "\n",
    "def plot_classes_preds(net, images, labels):\n",
    "    '''\n",
    "    Returns a plot using the network, along with images\n",
    "    and labels from a batch, that shows the network's class prediction along\n",
    "    with its probability, alongside the actual label, coloring this\n",
    "    information based on whether the prediction was correct or not.\n",
    "    Uses the \"images_to_probs\" function defined above.\n",
    "    '''\n",
    "    preds, probs = images_to_probs(net, images)\n",
    "    # plot the images in the batch, along with predicted and true labels\n",
    "    fig = plt.figure(figsize=(4,4))\n",
    "\n",
    "    for idx in np.arange(4):\n",
    "        ax = fig.add_subplot(4, 1, idx+1, xticks=[], yticks=[])\n",
    "        fig.tight_layout()\n",
    "        matplotlib_imshow(images[idx], one_channel=True)\n",
    "        ax.set_title(\"{0}, {1:.1f}%(label: {2})\".format(\n",
    "            classes[preds[idx]],\n",
    "            probs[idx] * 100.0,\n",
    "            classes[labels[idx]]),\n",
    "                    color=(\"green\" if preds[idx]==labels[idx].item() else \"red\"),loc=\"center\",pad=5,fontsize=\"medium\")\n",
    "    return fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "U_E9vMXRXlun"
   },
   "source": [
    "We are all set up to train the model! Let's use the same framework we used in the `PyTorch` tutorial notebook.\n",
    "\n",
    "Let's write the average loss and the plot generated from `plot_classes_preds` to TensorBoard every 1000 batches  using the `add_scalar` and `add_figure` functions.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lAS35sgKXlun"
   },
   "outputs": [],
   "source": [
    "epochs = 1\n",
    "running_loss = 0.0\n",
    "net.to(device)\n",
    "\n",
    "for epoch in range(epochs):  # loop over the dataset multiple times\n",
    "    for i, data in enumerate(trainloader, 0):#Iterating through the minibatches of the data\n",
    "\n",
    "        # data is a tuple of (inputs, labels)\n",
    "        inputs, labels = data\n",
    "\n",
    "        # Makes sure that the model and the data are in the same device\n",
    "        inputs = inputs.to(device)\n",
    "        labels = labels.to(device)\n",
    "\n",
    "        # Reset the parameter gradients for the current  minibatch iteration\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "\n",
    "        outputs = net(inputs)              # Perform a forward pass on the network with inputs\n",
    "        loss = criterion(outputs, labels)  # calculate the loss with the network predictions and ground Truth\n",
    "        loss.backward()                    # Perform a backward pass to calculate the gradients\n",
    "        optimizer.step()                   # Optimise the network parameters with calculated gradients\n",
    "\n",
    "        # Accumulate the loss\n",
    "        running_loss += loss.item()\n",
    "\n",
    "        if i % 1000 == 999:    # every thousandth mini-batch\n",
    "            print(\"[Epoch %d, Iteration %5d]\" % (epoch+1, i+1))\n",
    "\n",
    "            # log the running loss\n",
    "            writer.add_scalar('Training loss',\n",
    "                            running_loss / 1000,\n",
    "                            epoch * len(trainloader) + i)\n",
    "\n",
    "            # log the plot showing the model's predictions on a  sample of mini-batch using our helper function\n",
    "\n",
    "            writer.add_figure('Predictions vs Actuals',\n",
    "                            plot_classes_preds(net, inputs, labels),\n",
    "                            i)\n",
    "            running_loss = 0.0\n",
    "\n",
    "print('Finished Training')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UgE5zhsUXluo"
   },
   "source": [
    "You will now be able to see the plot of loss under `SCALARS` tab. We can also see the  figure for predicted samples in `IMAGES` tab."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zFdprQoLXluo"
   },
   "source": [
    "![](https://i2dl.vc.in.tum.de/static/images/exercise_07/tb_results1.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "T-oXCWeQXluo"
   },
   "source": [
    "![](https://i2dl.vc.in.tum.de/static/images/exercise_07/tb_results3.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e2ohgHeEXluo"
   },
   "source": [
    "# 6. Experimenting  weight initialization strategies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZrMIEDvlXlup"
   },
   "source": [
    "We will now  apply all the techniques we have learned in `TensorBoard` to explore the effect of different weight initializations. In the previous exercises, we used a naive Gaussian initialization, though in the lectures you learned that one needs to be careful about the weight initialization values. In addition, weight initialization is dependent on the activation function used.  \n",
    "\n",
    "Let's replicate those experiments!\n",
    "\n",
    "The code below initializes a new `SummaryWriter` instance to log experiment values in the directory `weight_init_experiments`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HE0I8qD9Xlup"
   },
   "outputs": [],
   "source": [
    "from torch.utils.tensorboard import SummaryWriter\n",
    "writer = SummaryWriter('logs/weight_init_experiments')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UA-jVRF5Xlup"
   },
   "source": [
    "Let's define a test network for the experiment and keep track of the output of each layer to find how the input data is modified through the layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gnuVOluZXlup"
   },
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class Net(nn.Module):\n",
    "    def __init__(self, activation_method):\n",
    "        super().__init__()\n",
    "\n",
    "        self.x1 = torch.Tensor([])\n",
    "        self.x2 = torch.Tensor([])\n",
    "        self.x3 = torch.Tensor([])\n",
    "        self.x4 = torch.Tensor([])\n",
    "        self.x5 = torch.Tensor([])\n",
    "        self.x6 = torch.Tensor([])\n",
    "\n",
    "        self.fc1 = nn.Linear(28*28, 300)\n",
    "        self.fc2 = nn.Linear(300, 300)\n",
    "        self.fc3 = nn.Linear(300, 300)\n",
    "        self.fc4 = nn.Linear(300, 300)\n",
    "        self.fc5 = nn.Linear(300, 300)\n",
    "        self.fc6 = nn.Linear(300, 300)\n",
    "        self.fc7 = nn.Linear(300, 10)\n",
    "\n",
    "        if activation_method == \"relu\" :\n",
    "            self.activation = nn.ReLU()\n",
    "        elif activation_method == \"tanh\":\n",
    "            self.activation = nn.Tanh()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.reshape(-1,28*28)\n",
    "        self.x1 = self.activation(self.fc1(x))\n",
    "        self.x2 = self.activation(self.fc2(self.x1))\n",
    "        self.x3 = self.activation(self.fc3(self.x2))\n",
    "        self.x4 = self.activation(self.fc4(self.x3))\n",
    "        self.x5 = self.activation(self.fc5(self.x4))\n",
    "        self.x6 = self.activation(self.fc6(self.x5))\n",
    "        logits = self.fc7 (self.x6)\n",
    "        return logits\n",
    "\n",
    "    def collect_layer_out (self):# Return the output values for each of the network layers\n",
    "        return [self.x1, self.x2, self.x3, self.x4, self.x5, self.x6]\n",
    "\n",
    "net = Net(\"tanh\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UQkG1y7qXlup"
   },
   "source": [
    "Let's now sample a batch of images for input to the network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xjddzNM0Xlup"
   },
   "outputs": [],
   "source": [
    "visloader = torch.utils.data.DataLoader(fashion_mnist_dataset, batch_size=40, shuffle=True)\n",
    "for images, labels in visloader:\n",
    "    break\n",
    "\n",
    "print(\"Size of the Mini-batch input:\",images.size())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GbNVblzGXluq"
   },
   "source": [
    "We will plot the histogram of activation values  produced in each of the network layers as the input passes through the network model using the `add_histogram` function. This helps us look at the distribution of activation values. Select the `HISTOGRAMS` tab in TensorBoard to visualise the experiment results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6i9D1IhMXluq"
   },
   "source": [
    "## 6.1 Constant weight initialization with $tanh$ activation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Sqe6I1tJXluq"
   },
   "source": [
    "Let's start with constant weight initialization. What problems do you observe with the distribution of the output of each layer?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5MQSqd3wXluq"
   },
   "outputs": [],
   "source": [
    "net_const = Net(\"tanh\")\n",
    "\n",
    "def init_weights(m):\n",
    "    if type(m) == nn.Linear:\n",
    "        torch.nn.init.constant_(m.weight,2.0)\n",
    "        m.bias.data.fill_(0.01)\n",
    "\n",
    "net_const.apply(init_weights)\n",
    "outputs = net_const(images)\n",
    "layer_out = net_const.collect_layer_out()\n",
    "\n",
    "for i, x in enumerate(layer_out):\n",
    "    writer.add_histogram('constant_init', x, i+1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7KRcRyIIXluq"
   },
   "source": [
    "We can see that initialization with constant values does not break the symmetry of weights, i.e. all neurons in network always learn the same features from the input since the weights are the same.  \n",
    "\n",
    "Now we will try random weight initialization and let's see what happens if weights are initialized with high numerical values or very low numerical values."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HQFtNu0zXluq"
   },
   "source": [
    "## 6.2 Random weight initialization of small numerical values with $tanh$ activation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gDY2ilHtXlur"
   },
   "outputs": [],
   "source": [
    "net_small_normal = Net(\"tanh\")\n",
    "\n",
    "def init_weights(m):\n",
    "    if type(m) == nn.Linear:\n",
    "        torch.nn.init.normal_(m.weight,mean=0.0, std=0.01)\n",
    "        m.bias.data.fill_(0.01)\n",
    "\n",
    "net_small_normal.apply(init_weights)\n",
    "outputs = net_small_normal(images)\n",
    "layer_out = net_small_normal.collect_layer_out()\n",
    "\n",
    "for i, x in enumerate(layer_out):\n",
    "    writer.add_histogram('small_normal_tanh', x, i+1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sJpBfFmcXlur"
   },
   "source": [
    "## 6.3  Random weight initialization of large numerical values with $tanh$ activation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Hm4Gj4tQXlur"
   },
   "outputs": [],
   "source": [
    "net_large_normal = Net(\"tanh\")\n",
    "\n",
    "def init_weights(m):\n",
    "    if type(m) == nn.Linear:\n",
    "        torch.nn.init.normal_(m.weight,mean=0.0, std=0.2)\n",
    "        m.bias.data.fill_(0.01)\n",
    "\n",
    "net_large_normal.apply(init_weights)\n",
    "outputs = net_large_normal(images)\n",
    "layer_out = net_large_normal.collect_layer_out()\n",
    "\n",
    "for i, x in enumerate(layer_out):\n",
    "    writer.add_histogram('large_normal_tanh', x, i+1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "q-t3tEJgXlur"
   },
   "source": [
    "From last two examples, we can see that random weight initialization with normal distribution might work well in some shallow layers of the network, while if we are going deeper into the network, it will end up with **vanishing gradient problem**, i.e.\n",
    "\n",
    "- If weights are initialized with very high values, the term $Xw+b$ becomes significantly higher and with activation function such as $tanh$, the function returns value very close to $-1$ or $1$. At these values, the gradient of $tanh$ is very low, thus learning takes a lot of time.\n",
    "\n",
    "- If weights are initialized with low values, it gets mapped to around 0, and the small values will kill gradients when backpropagating through the network.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "trNZ54KEXlur"
   },
   "source": [
    "## 6.4 Xavier initialization with $tanh$ activation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-CB-ZQypXlur"
   },
   "source": [
    "From the previous examples, we can see that a proper weight initialization is needed to ensure nice distribution of the output of each layer. Here comes the **Xavier Initialization**.\n",
    "\n",
    "We will fill the weight with values using a normal distribution $\\mathcal{N}(0,{\\sigma}^2)$ where\n",
    "\n",
    "$$ \\sigma = gain \\times \\sqrt{\\frac{2}{fan _{in} + fan_{out}}} $$\n",
    "\n",
    "Here $fan _{in}$ and $ fan_{out} $ are number of neurons in the input and output layer and ${gain}$ is a optional scaling factor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YscO2pLLXlus"
   },
   "outputs": [],
   "source": [
    "net_xavier = Net(\"tanh\")\n",
    "\n",
    "def init_weights(m):\n",
    "    if type(m) == nn.Linear:\n",
    "        torch.nn.init.xavier_normal_(m.weight)\n",
    "        m.bias.data.fill_(0.01)\n",
    "\n",
    "net_xavier.apply(init_weights)\n",
    "outputs = net_xavier(images)\n",
    "layer_out = net_xavier.collect_layer_out()\n",
    "\n",
    "for i, x in enumerate(layer_out):\n",
    "    writer.add_histogram('xavier_tanh', x, i+1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8fwLXFQ-Xlus"
   },
   "source": [
    "## 6.5 Xavier initialization with ReLU"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rbdB8Jh0Xlus"
   },
   "source": [
    "Xavier initialization requires a zero centered activation function such as $tanh$ to work well. Let's try using the Xavier initialization with ReLU:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2SvAHe4zXlus"
   },
   "outputs": [],
   "source": [
    "net_xavier_relu = Net(\"relu\")\n",
    "\n",
    "def init_weights(m):\n",
    "    if type(m) == nn.Linear:\n",
    "        torch.nn.init.xavier_uniform_(m.weight)\n",
    "        m.bias.data.fill_(0.01)\n",
    "\n",
    "net_xavier_relu.apply(init_weights)\n",
    "outputs = net_xavier_relu(images)\n",
    "layer_out = net_xavier_relu.collect_layer_out()\n",
    "\n",
    "for i, x in enumerate(layer_out):\n",
    "    writer.add_histogram('xavier_relu', x, i+1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Gm06ncjtXlus"
   },
   "source": [
    "We can see here that layer outputs collapse to zero again if we use non-zero centered activation such as ReLU."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WygBMr6eXlus"
   },
   "source": [
    "## 6.6 He initialization with ReLU"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IWr6leV5Xlus"
   },
   "source": [
    "**He Initialization** comes to our rescue for non-centered activation functions. We will fill the weight with values using a normal distribution $\\mathcal{N}(0,\\sigma^2)$ where\n",
    "\n",
    "$$ \\sigma = \\frac {gain} {\\sqrt{fan_{mode}}} $$\n",
    "\n",
    "Here $fan _{mode}$ can be chosen either $fan _{in}$ (default) or $fan _{out}$.\n",
    "\n",
    "Choosing $fan _{in}$ preserves the magnitude of the variance of the weights in the forward pass. Choosing $fan _{out}$ preserves the magnitudes of weights during the backwards pass. The variable $gain$ is again the optional scaling factor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oaObyjItXlus"
   },
   "outputs": [],
   "source": [
    "net_kaiming_relu = Net(\"relu\")\n",
    "\n",
    "def init_weights(m):\n",
    "    if type(m) == nn.Linear:\n",
    "        torch.nn.init.kaiming_uniform_(m.weight,nonlinearity='relu')\n",
    "        m.bias.data.fill_(0.01)\n",
    "\n",
    "net_kaiming_relu.apply(init_weights)\n",
    "outputs = net_kaiming_relu(images)\n",
    "layer_out = net_kaiming_relu.collect_layer_out()\n",
    "\n",
    "for i, x in enumerate(layer_out):\n",
    "    writer.add_histogram('kaiming_relu', x, i+1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WG9jZC3fXlut"
   },
   "source": [
    "With these, you should have everything at hand to work with Tensorboard. It is highly advised to use either Tensorboard or other similar libraries, such as visdom to visualise network training results.\n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "i2dl",
   "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.11.11"
  },
  "vscode": {
   "interpreter": {
    "hash": "54970da6898dad277dbf355945c2dee7f942d2a31ec1fc1455b6d4f552d07b83"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
