0_illustration_SoftmaxCosineSim.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Illustration for custom layer: Softmax_cosine_sim\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "from random import shuffle\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras.activations import softmax \n",
    "\n",
    "class SoftmaxCosineSim(keras.layers.Layer):\n",
    "    # ==============================================================================\n",
    "    # Code modified from NT-XENT-loss: \n",
    "    # https://github.com/google-research/simclr/blob/master/objective.py\n",
    "    # ==============================================================================\n",
    "    # coding=utf-8\n",
    "    # Copyright 2020 The SimCLR Authors.\n",
    "    #\n",
    "    # Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "    # you may not use this file except in compliance with the License.\n",
    "    # You may obtain a copy of the License at\n",
    "    #\n",
    "    #     http://www.apache.org/licenses/LICENSE-2.0\n",
    "    #\n",
    "    # Unless required by applicable law or agreed to in writing, software\n",
    "    # distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "    # See the License for the specific simclr governing permissions and\n",
    "    # limitations under the License.\n",
    "    # ==============================================================================\n",
    "    def __init__(self, batch_size, feat_dim, temperature = 0.1, LARGE_NUM = 1e9):\n",
    "        self.batch_size = batch_size\n",
    "        self.feat_dim = feat_dim\n",
    "        self.units = (batch_size, 4 * feat_dim)\n",
    "        self.input_dim = [(None, feat_dim)] * (batch_size * 2)\n",
    "        self.temperature = temperature\n",
    "        self.LARGE_NUM = LARGE_NUM\n",
    "        super(SoftmaxCosineSim, self).__init__()\n",
    "        \n",
    "    def get_config(self):\n",
    "\n",
    "        config = super().get_config().copy()\n",
    "        config.update({\n",
    "            'batch_size' : self.batch_size,\n",
    "            'feat_dim' : self.feat_dim,\n",
    "            'units' : self.units,\n",
    "            'input_dim' : self.input_dim,\n",
    "            'temperature' : self.temperature,\n",
    "            'LARGE_NUM' : self.LARGE_NUM,\n",
    "        })\n",
    "        return config    \n",
    "\n",
    "    def call(self, inputs):\n",
    "        # Function to perform tranformatiom: input -> output\n",
    "        batch_size = len(inputs) // 2\n",
    "        z1 = []\n",
    "        z2 = []\n",
    "        \n",
    "        for index in range(batch_size):\n",
    "            # 0 assumes that batch in generator is actually just 1\n",
    "            z1.append(tf.math.l2_normalize(inputs[index][0], -1))\n",
    "            z2.append(tf.math.l2_normalize(inputs[batch_size + index][0], -1))\n",
    "        \n",
    "        # Gather hidden1/hidden2 across replicas and create local labels.\n",
    "        z1_large = z1\n",
    "        z2_large = z2\n",
    "                      \n",
    "        # TODO: move to GENERATOR \n",
    "        labels = tf.one_hot(tf.range(batch_size), batch_size * 2)\n",
    "        masks = tf.one_hot(tf.range(batch_size), batch_size)\n",
    "\n",
    "        # Products of vectors of same side of network (z_i), count as negative examples\n",
    "        logits_aa = tf.matmul(z1, z1_large, transpose_b=True) / self.temperature\n",
    "        print(f\"logits_aa \\n {logits_aa} \\n\")\n",
    "\n",
    "        # Values on the diagonal are put equal to a very small value -> exclude product between 2 identical values\n",
    "        logits_aa = logits_aa - masks * self.LARGE_NUM\n",
    "        print(f\"logits_aa after mask:: \\n {logits_aa} \\n\")\n",
    "\n",
    "        # Similar as aa\n",
    "        logits_bb = tf.matmul(z2, z2_large, transpose_b=True) / self.temperature\n",
    "        logits_bb = logits_bb - masks * self.LARGE_NUM\n",
    "\n",
    "\n",
    "        # Comparison between two sides of the network (z_i and z_j) -> diagonal should be as close as possible to 1\n",
    "        logits_ab = tf.matmul(z1, z2_large, transpose_b=True) / self.temperature\n",
    "        print(f\"logits_ab: \\n {logits_ab} \\n\")\n",
    "        # Similar as ba\n",
    "        logits_ba = tf.matmul(z2, z1_large, transpose_b=True) / self.temperature\n",
    "        \n",
    "        print(tf.concat([logits_ab, logits_aa], 1))\n",
    "        \n",
    "        part1 = softmax(tf.concat([logits_ab, logits_aa], 1))\n",
    "        part2 = softmax(tf.concat([logits_ba, logits_bb], 1))\n",
    "        output = tf.concat([part1, part2], 1)\n",
    "        \n",
    "        return output \n",
    "\n",
    "def tf_round_decimal(x, decimals = 0):\n",
    "    multiplier = tf.constant(10**decimals, dtype=x.dtype)\n",
    "    return tf.round(x * multiplier) / multiplier"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Strictly ordered feature vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 3\n",
    "feat_dim = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "hidden_ordered = [tf.convert_to_tensor(np.asarray([[1,0,0,0,0]]).astype('float32')), \n",
    "          tf.convert_to_tensor(np.asarray([[0,0,0,0,1]]).astype('float32')),\n",
    "          tf.convert_to_tensor(np.asarray([[0,0,1,0,0]]).astype('float32')),\n",
    "          tf.convert_to_tensor(np.asarray([[1,0,0,0,0]]).astype('float32')), \n",
    "          tf.convert_to_tensor(np.asarray([[0,0,0,0,1]]).astype('float32')), \n",
    "          tf.convert_to_tensor(np.asarray([[0,0,1,0,0]]).astype('float32'))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "logits_aa \n",
      " [[10.  0.  0.]\n",
      " [ 0. 10.  0.]\n",
      " [ 0.  0. 10.]] \n",
      "\n",
      "logits_aa after mask:: \n",
      " [[-1.e+09  0.e+00  0.e+00]\n",
      " [ 0.e+00 -1.e+09  0.e+00]\n",
      " [ 0.e+00  0.e+00 -1.e+09]] \n",
      "\n",
      "logits_ab: \n",
      " [[10.  0.  0.]\n",
      " [ 0. 10.  0.]\n",
      " [ 0.  0. 10.]] \n",
      "\n",
      "tf.Tensor(\n",
      "[[ 1.e+01  0.e+00  0.e+00 -1.e+09  0.e+00  0.e+00]\n",
      " [ 0.e+00  1.e+01  0.e+00  0.e+00 -1.e+09  0.e+00]\n",
      " [ 0.e+00  0.e+00  1.e+01  0.e+00  0.e+00 -1.e+09]], shape=(3, 6), dtype=float32)\n",
      " \n",
      " output: \n",
      " \n",
      " [[0.9998 0.     0.     0.     0.     0.     0.9998 0.     0.     0.\n",
      "  0.     0.    ]\n",
      " [0.     0.9998 0.     0.     0.     0.     0.     0.9998 0.     0.\n",
      "  0.     0.    ]\n",
      " [0.     0.     0.9998 0.     0.     0.     0.     0.     0.9998 0.\n",
      "  0.     0.    ]]\n"
     ]
    }
   ],
   "source": [
    "SoftmaxCosineSim_layer = SoftmaxCosineSim(batch_size = batch_size, feat_dim = feat_dim)\n",
    "y = SoftmaxCosineSim_layer(hidden_ordered)\n",
    "print(f\" \\n output: \\n \\n {tf_round_decimal(y,4)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Shuffled feature vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1. 0. 0. 0. 0.]\n",
      "[0. 0. 0. 0. 1.]\n",
      "[0. 0. 1. 0. 0.]\n",
      "[0. 0. 1. 0. 0.]\n",
      "[1. 0. 0. 0. 0.]\n",
      "[0. 0. 0. 0. 1.]\n"
     ]
    }
   ],
   "source": [
    "hidden_shuffled = [tf.convert_to_tensor(np.asarray([[1,0,0,0,0]]).astype('float32')), \n",
    "                   tf.convert_to_tensor(np.asarray([[0,0,0,0,1]]).astype('float32')),\n",
    "                   tf.convert_to_tensor(np.asarray([[0,0,1,0,0]]).astype('float32')),\n",
    "                   tf.convert_to_tensor(np.asarray([[0,0,1,0,0]]).astype('float32')),\n",
    "                   tf.convert_to_tensor(np.asarray([[1,0,0,0,0]]).astype('float32')),\n",
    "                   tf.convert_to_tensor(np.asarray([[0,0,0,0,1]]).astype('float32'))]\n",
    "for h in hidden_shuffled:\n",
    "    print(h.numpy()[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### => This should result on 1 on index (1,2), (2,3), (3,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "logits_aa \n",
      " [[10.  0.  0.]\n",
      " [ 0. 10.  0.]\n",
      " [ 0.  0. 10.]] \n",
      "\n",
      "logits_aa after mask:: \n",
      " [[-1.e+09  0.e+00  0.e+00]\n",
      " [ 0.e+00 -1.e+09  0.e+00]\n",
      " [ 0.e+00  0.e+00 -1.e+09]] \n",
      "\n",
      "logits_ab: \n",
      " [[ 0. 10.  0.]\n",
      " [ 0.  0. 10.]\n",
      " [10.  0.  0.]] \n",
      "\n",
      "tf.Tensor(\n",
      "[[ 0.e+00  1.e+01  0.e+00 -1.e+09  0.e+00  0.e+00]\n",
      " [ 0.e+00  0.e+00  1.e+01  0.e+00 -1.e+09  0.e+00]\n",
      " [ 1.e+01  0.e+00  0.e+00  0.e+00  0.e+00 -1.e+09]], shape=(3, 6), dtype=float32)\n",
      " \n",
      " output: \n",
      " \n",
      " [[0.     0.9998 0.     0.     0.     0.     0.     0.     0.9998 0.\n",
      "  0.     0.    ]\n",
      " [0.     0.     0.9998 0.     0.     0.     0.9998 0.     0.     0.\n",
      "  0.     0.    ]\n",
      " [0.9998 0.     0.     0.     0.     0.     0.     0.9998 0.     0.\n",
      "  0.     0.    ]]\n"
     ]
    }
   ],
   "source": [
    "y = SoftmaxCosineSim_layer(hidden_shuffled)\n",
    "print(f\" \\n output: \\n \\n {tf_round_decimal(y,4)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### => Indeed, 1 on index (1,2), (2,3), (3,1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Illustration of temperature importance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "logits_aa \n",
      " [[1. 0. 0.]\n",
      " [0. 1. 0.]\n",
      " [0. 0. 1.]] \n",
      "\n",
      "logits_aa after mask:: \n",
      " [[-1.e+09  0.e+00  0.e+00]\n",
      " [ 0.e+00 -1.e+09  0.e+00]\n",
      " [ 0.e+00  0.e+00 -1.e+09]] \n",
      "\n",
      "logits_ab: \n",
      " [[1. 0. 0.]\n",
      " [0. 1. 0.]\n",
      " [0. 0. 1.]] \n",
      "\n",
      "tf.Tensor(\n",
      "[[ 1.e+00  0.e+00  0.e+00 -1.e+09  0.e+00  0.e+00]\n",
      " [ 0.e+00  1.e+00  0.e+00  0.e+00 -1.e+09  0.e+00]\n",
      " [ 0.e+00  0.e+00  1.e+00  0.e+00  0.e+00 -1.e+09]], shape=(3, 6), dtype=float32)\n",
      " \n",
      " output: \n",
      " \n",
      " [[0.405 0.149 0.149 0.    0.149 0.149 0.405 0.149 0.149 0.    0.149 0.149]\n",
      " [0.149 0.405 0.149 0.149 0.    0.149 0.149 0.405 0.149 0.149 0.    0.149]\n",
      " [0.149 0.149 0.405 0.149 0.149 0.    0.149 0.149 0.405 0.149 0.149 0.   ]]\n"
     ]
    }
   ],
   "source": [
    "SoftmaxCosineSim_layer = SoftmaxCosineSim(batch_size = batch_size, feat_dim = feat_dim, temperature = 1)\n",
    "y = SoftmaxCosineSim_layer(hidden_ordered)\n",
    "print(f\" \\n output: \\n \\n {tf_round_decimal(y,3)}\")"
   ]
  },
  {
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}