BatchNormalization.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.1.0'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Batch Normalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 3), dtype=float32, numpy=\n",
       "array([[1., 2., 3.],\n",
       "       [4., 5., 6.]], dtype=float32)>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = tf.reshape(tf.range(1,7, dtype = tf.float32), shape=(2,3))\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class custom_bn(tf.keras.layers.Layer):\n",
    "    \n",
    "    def __init__(self, scale=1, offset=0, axes=[0], **kwargs):\n",
    "        super(custom_bn, self).__init__(**kwargs)\n",
    "        \"\"\"\n",
    "        A simple batch normalization layer. \n",
    "        Check Keras BatchNormalization for advanced customization\n",
    "        \n",
    "        Also, the $/sigma$ in the equation in the tf API referring to variance\n",
    "        is the standard deviation of the variance, not the variance itself\n",
    "        \n",
    "        Parameters\n",
    "        ==========\n",
    "        axes: list; passed as argument to tf.nn.moments for calculating the mean and variance\n",
    "        \"\"\"\n",
    "        \n",
    "        self.scale = scale\n",
    "        self.offset = offset\n",
    "        self.axes = axes\n",
    "        \n",
    "    def batch_norm_fn(self, inputs):\n",
    "        \n",
    "        self.mean, self.variance = tf.nn.moments(inputs, self.axes, keepdims=True)\n",
    "        # If keepdims = False, resulting in (for axes = [1])\n",
    "        # \"InvalidArgumentError: Incompatible shapes: [2,3] vs. [2] [Op:Mul] name: custom_bn_1/cond/batchnorm/mul/\"\n",
    "        # This is because the output shapes returned from tf.nn.moments are not compatible for broadcasting\n",
    "        # depending on which axes is used. \n",
    "        # Changing keepdims = True solves this   \n",
    "\n",
    "        \n",
    "        \n",
    "        return tf.nn.batch_normalization(inputs, mean = self.mean, variance= self.variance,\n",
    "                                                                   scale = self.scale, offset = self.offset, \n",
    "                                                                  variance_epsilon=0.001)\n",
    "        \n",
    "    def call(self, inputs, training = None):\n",
    "        # If proper arguments are not passed to the batch_norm function below, then the following error will be output\n",
    "        # \"TypeError: batch_normalization() missing 5 required positional arguments: \n",
    "        #    'mean', 'variance', 'offset', 'scale', and 'variance_epsilon'\"\n",
    "        \n",
    "        return tf.cond(training, lambda: self.batch_norm_fn(inputs), lambda: inputs)\n",
    "            "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Testing on 2D data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Axis = 0\n",
    "custom_bn_layer = custom_bn(scale =2, offset = 1, axes=[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 3), dtype=float32, numpy=\n",
       "array([[1., 2., 3.],\n",
       "       [4., 5., 6.]], dtype=float32)>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "custom_bn_layer(a) # defaulting training = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 3), dtype=float32, numpy=\n",
       "array([[-0.9995558, -0.9995556, -0.9995558],\n",
       "       [ 2.9995556,  2.999556 ,  2.9995556]], dtype=float32)>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "custom_bn_layer(a, training = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Axis = 1\n",
    "custom_bn_layer2 = custom_bn(scale =2, offset = 1, axes=[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 3), dtype=float32, numpy=\n",
       "array([[-1.4476547,  1.       ,  3.4476547],\n",
       "       [-1.4476547,  1.       ,  3.4476547]], dtype=float32)>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "custom_bn_layer2(a, training = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### *Testing: Dismantling the batch normalization process*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 3), dtype=float32, numpy=\n",
       "array([[1., 2., 3.],\n",
       "       [4., 5., 6.]], dtype=float32)>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor([2.5 3.5 4.5], shape=(3,), dtype=float32) tf.Tensor([2.25 2.25 2.25], shape=(3,), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "mean, variance = tf.nn.moments(a, axes = [0], keepdims=False)\n",
    "print(mean, variance)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<tf.Tensor: shape=(1, 3), dtype=float32, numpy=array([[2.5, 3.5, 4.5]], dtype=float32)>,\n",
       " <tf.Tensor: shape=(1, 3), dtype=float32, numpy=array([[2.25, 2.25, 2.25]], dtype=float32)>)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.nn.moments(a, axes = [0], keepdims=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "    * Axes = [0] calculates mean across the batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 3), dtype=float32, numpy=\n",
       "array([[-1.5, -1.5, -1.5],\n",
       "       [ 1.5,  1.5,  1.5]], dtype=float32)>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a - mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 3), dtype=float32, numpy=\n",
       "array([[-1.5, -1.5, -1.5],\n",
       "       [ 1.5,  1.5,  1.5]], dtype=float32)>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a - tf.nn.moments(a, axes = [0], keepdims=True)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "scale, offset = 2, 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 3), dtype=float32, numpy=\n",
       "array([[-3., -3., -3.],\n",
       "       [ 3.,  3.,  3.]], dtype=float32)>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(a-mean)*scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 3), dtype=float32, numpy=\n",
       "array([[-0.9995557, -0.9995557, -0.9995557],\n",
       "       [ 2.9995556,  2.9995556,  2.9995556]], dtype=float32)>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(a-mean)/(tf.math.sqrt(variance + 0.001))*scale + offset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 3), dtype=float32, numpy=\n",
       "array([[-0.9995557, -0.9995557, -0.9995557],\n",
       "       [ 2.9995556,  2.9995556,  2.9995556]], dtype=float32)>"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# using tf functions\n",
    "tf.add(tf.multiply(tf.divide(tf.subtract(a, mean), tf.math.sqrt(tf.add(variance, 0.001))), scale), offset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "    * Axes = [1] should calculate the moments along the columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[2.]\n",
      " [5.]], shape=(2, 1), dtype=float32) tf.Tensor(\n",
      "[[0.6666667]\n",
      " [0.6666667]], shape=(2, 1), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "mean, variance = tf.nn.moments(a, axes = [1], keepdims=True)\n",
    "print(mean, variance)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 3), dtype=float32, numpy=\n",
       "array([[1., 2., 3.],\n",
       "       [4., 5., 6.]], dtype=float32)>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([2., 5.], dtype=float32)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a.numpy().mean(axis = 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.6666667, 0.6666667], dtype=float32)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a.numpy().var(axis = 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 3), dtype=float32, numpy=\n",
       "array([[-1.4476547,  1.       ,  3.4476547],\n",
       "       [-1.4476547,  1.       ,  3.4476547]], dtype=float32)>"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# using tf functions\n",
    "tf.add(tf.multiply(tf.divide(tf.subtract(a, mean), tf.math.sqrt(tf.add(variance, 0.001))), scale), offset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "    * Both axes = [0] and [1] confirms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Testing on 4D data such as in images or timeseries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2, 3, 2), dtype=float32, numpy=\n",
       "array([[[[ 0.,  1.],\n",
       "         [ 2.,  3.],\n",
       "         [ 4.,  5.]],\n",
       "\n",
       "        [[ 6.,  7.],\n",
       "         [ 8.,  9.],\n",
       "         [10., 11.]]],\n",
       "\n",
       "\n",
       "       [[[12., 13.],\n",
       "         [14., 15.],\n",
       "         [16., 17.]],\n",
       "\n",
       "        [[18., 19.],\n",
       "         [20., 21.],\n",
       "         [22., 23.]]]], dtype=float32)>"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_data = tf.reshape(tf.range(0,24, dtype = np.float32), (2,2,3,2)) # reshaping in the format of an image [batch, height, width, channels]\n",
    "test_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Axes = [0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2, 3, 2), dtype=float32, numpy=\n",
       "array([[[[-0.9999721 , -0.9999721 ],\n",
       "         [-0.99997216, -0.9999722 ],\n",
       "         [-0.9999722 , -0.9999722 ]],\n",
       "\n",
       "        [[-0.9999721 , -0.9999721 ],\n",
       "         [-0.9999721 , -0.9999721 ],\n",
       "         [-0.9999721 , -0.9999721 ]]],\n",
       "\n",
       "\n",
       "       [[[ 2.999972  ,  2.999972  ],\n",
       "         [ 2.999972  ,  2.999972  ],\n",
       "         [ 2.999972  ,  2.999972  ]],\n",
       "\n",
       "        [[ 2.9999723 ,  2.9999723 ],\n",
       "         [ 2.9999723 ,  2.9999723 ],\n",
       "         [ 2.9999723 ,  2.9999723 ]]]], dtype=float32)>"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "custom_bn_layer = custom_bn(scale=2, offset=1, axes=[0])\n",
    "x = custom_bn_layer(test_data, training = True)\n",
    "x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### *Testing: Dismantling the batch normalization process*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[[[ 6.  7.]\n",
      "   [ 8.  9.]\n",
      "   [10. 11.]]\n",
      "\n",
      "  [[12. 13.]\n",
      "   [14. 15.]\n",
      "   [16. 17.]]]], shape=(1, 2, 3, 2), dtype=float32) tf.Tensor(\n",
      "[[[[36. 36.]\n",
      "   [36. 36.]\n",
      "   [36. 36.]]\n",
      "\n",
      "  [[36. 36.]\n",
      "   [36. 36.]\n",
      "   [36. 36.]]]], shape=(1, 2, 3, 2), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# mean, var should have the shape [1,d2,d3,d4] where d2,d3,d4 are respective dimensions of the 4D input\n",
    "mean, variance = tf.nn.moments(test_data, axes = [0], keepdims=True)\n",
    "print(mean, variance)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2, 3, 2), dtype=float32, numpy=\n",
       "array([[[[-1.0000002, -1.0000002],\n",
       "         [-1.0000002, -1.0000002],\n",
       "         [-1.0000002, -1.0000002]],\n",
       "\n",
       "        [[-1.0000002, -1.0000002],\n",
       "         [-1.       , -1.       ],\n",
       "         [-1.       , -1.       ]]],\n",
       "\n",
       "\n",
       "       [[[ 3.0000002,  3.0000002],\n",
       "         [ 3.0000002,  3.0000002],\n",
       "         [ 3.0000002,  3.0000002]],\n",
       "\n",
       "        [[ 3.0000002,  3.0000002],\n",
       "         [ 3.       ,  3.       ],\n",
       "         [ 3.       ,  3.       ]]]], dtype=float32)>"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# using tf functions\n",
    "xhat = tf.add(tf.multiply(tf.divide(tf.subtract(test_data, mean), tf.math.sqrt(variance, 0.001)), scale), offset)\n",
    "xhat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2, 3, 2), dtype=float32, numpy=\n",
       "array([[[[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]],\n",
       "\n",
       "        [[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]]],\n",
       "\n",
       "\n",
       "       [[[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]],\n",
       "\n",
       "        [[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]]]], dtype=float32)>"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.round(x-xhat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "    * Confirms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Axes = [1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2, 3, 2), dtype=float32, numpy=\n",
       "array([[[[-0.9998889 , -0.99988884],\n",
       "         [-0.9998888 , -0.9998889 ],\n",
       "         [-0.99988866, -0.9998889 ]],\n",
       "\n",
       "        [[ 2.999889  ,  2.9998887 ],\n",
       "         [ 2.999889  ,  2.999889  ],\n",
       "         [ 2.999889  ,  2.999889  ]]],\n",
       "\n",
       "\n",
       "       [[[-0.9998884 , -0.9998884 ],\n",
       "         [-0.9998894 , -0.9998894 ],\n",
       "         [-0.9998884 , -0.9998884 ]],\n",
       "\n",
       "        [[ 2.9998894 ,  2.9998884 ],\n",
       "         [ 2.9998884 ,  2.9998884 ],\n",
       "         [ 2.9998894 ,  2.9998894 ]]]], dtype=float32)>"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "custom_bn_layer = custom_bn(scale=2, offset=1, axes=[1])\n",
    "x = custom_bn_layer(test_data, training = True)\n",
    "x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### *Testing: Dismantling the batch normalization process*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[[[ 3.  4.]\n",
      "   [ 5.  6.]\n",
      "   [ 7.  8.]]]\n",
      "\n",
      "\n",
      " [[[15. 16.]\n",
      "   [17. 18.]\n",
      "   [19. 20.]]]], shape=(2, 1, 3, 2), dtype=float32) tf.Tensor(\n",
      "[[[[9. 9.]\n",
      "   [9. 9.]\n",
      "   [9. 9.]]]\n",
      "\n",
      "\n",
      " [[[9. 9.]\n",
      "   [9. 9.]\n",
      "   [9. 9.]]]], shape=(2, 1, 3, 2), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# mean, var should have the shape [d1,1,d3,d4] where d1,d3,d4 are respective dimensions of the 4D input\n",
    "mean, variance = tf.nn.moments(test_data, axes = [1], keepdims=True)\n",
    "print(mean, variance)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2, 3, 2), dtype=float32, numpy=\n",
       "array([[[[-0.9998889, -0.9998889],\n",
       "         [-0.9998889, -0.9998889],\n",
       "         [-0.9998889, -0.9998889]],\n",
       "\n",
       "        [[ 2.999889 ,  2.999889 ],\n",
       "         [ 2.999889 ,  2.999889 ],\n",
       "         [ 2.999889 ,  2.999889 ]]],\n",
       "\n",
       "\n",
       "       [[[-0.9998889, -0.9998889],\n",
       "         [-0.9998889, -0.9998889],\n",
       "         [-0.9998889, -0.9998889]],\n",
       "\n",
       "        [[ 2.999889 ,  2.999889 ],\n",
       "         [ 2.999889 ,  2.999889 ],\n",
       "         [ 2.999889 ,  2.999889 ]]]], dtype=float32)>"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# using tf functions\n",
    "xhat = tf.add(tf.multiply(tf.divide(tf.subtract(test_data, mean), tf.math.sqrt(tf.add(variance, 0.001))), scale), offset)\n",
    "xhat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2, 3, 2), dtype=float32, numpy=\n",
       "array([[[[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]],\n",
       "\n",
       "        [[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]]],\n",
       "\n",
       "\n",
       "       [[[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]],\n",
       "\n",
       "        [[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]]]], dtype=float32)>"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.round(x-xhat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "    * Confirms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Axes = [2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2, 3, 2), dtype=float32, numpy=\n",
       "array([[[[-1.4490306, -1.4490305],\n",
       "         [ 1.       ,  1.       ],\n",
       "         [ 3.4490306,  3.4490309]],\n",
       "\n",
       "        [[-1.4490309, -1.4490299],\n",
       "         [ 1.       ,  1.       ],\n",
       "         [ 3.4490309,  3.4490309]]],\n",
       "\n",
       "\n",
       "       [[[-1.4490318, -1.4490299],\n",
       "         [ 1.       ,  1.       ],\n",
       "         [ 3.44903  ,  3.4490318]],\n",
       "\n",
       "        [[-1.4490318, -1.4490299],\n",
       "         [ 1.       ,  1.       ],\n",
       "         [ 3.44903  ,  3.4490318]]]], dtype=float32)>"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "custom_bn_layer = custom_bn(scale=2, offset=1, axes=[2])\n",
    "x = custom_bn_layer(test_data, training = True)\n",
    "x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### *Testing: Dismantling the batch normalization process*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[[[ 2.  3.]]\n",
      "\n",
      "  [[ 8.  9.]]]\n",
      "\n",
      "\n",
      " [[[14. 15.]]\n",
      "\n",
      "  [[20. 21.]]]], shape=(2, 2, 1, 2), dtype=float32) tf.Tensor(\n",
      "[[[[2.6666667 2.6666667]]\n",
      "\n",
      "  [[2.6666667 2.6666667]]]\n",
      "\n",
      "\n",
      " [[[2.6666667 2.6666667]]\n",
      "\n",
      "  [[2.6666667 2.6666667]]]], shape=(2, 2, 1, 2), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# mean, var should have the shape [d1,d2,1,d4] where d1,d2,d4 are respective dimensions of the 4D input\n",
    "mean, variance = tf.nn.moments(test_data, axes = [2], keepdims=True)\n",
    "print(mean, variance)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2, 3, 2), dtype=float32, numpy=\n",
       "array([[[[-1.4490306, -1.4490306],\n",
       "         [ 1.       ,  1.       ],\n",
       "         [ 3.4490306,  3.4490306]],\n",
       "\n",
       "        [[-1.4490306, -1.4490306],\n",
       "         [ 1.       ,  1.       ],\n",
       "         [ 3.4490306,  3.4490306]]],\n",
       "\n",
       "\n",
       "       [[[-1.4490306, -1.4490306],\n",
       "         [ 1.       ,  1.       ],\n",
       "         [ 3.4490306,  3.4490306]],\n",
       "\n",
       "        [[-1.4490306, -1.4490306],\n",
       "         [ 1.       ,  1.       ],\n",
       "         [ 3.4490306,  3.4490306]]]], dtype=float32)>"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# using tf functions\n",
    "xhat = tf.add(tf.multiply(tf.divide(tf.subtract(test_data, mean), tf.math.sqrt(tf.add(variance, 0.001))), scale), offset)\n",
    "xhat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2, 3, 2), dtype=float32, numpy=\n",
       "array([[[[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]],\n",
       "\n",
       "        [[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]]],\n",
       "\n",
       "\n",
       "       [[[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]],\n",
       "\n",
       "        [[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]]]], dtype=float32)>"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.round(x-xhat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "    * Confirms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Axes = [3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2, 3, 2), dtype=float32, numpy=\n",
       "array([[[[-0.99601173,  2.9960117 ],\n",
       "         [-0.99601173,  2.9960117 ],\n",
       "         [-0.99601173,  2.9960117 ]],\n",
       "\n",
       "        [[-0.99601173,  2.9960117 ],\n",
       "         [-0.9960098 ,  2.9960136 ],\n",
       "         [-0.9960098 ,  2.9960136 ]]],\n",
       "\n",
       "\n",
       "       [[[-0.9960098 ,  2.9960136 ],\n",
       "         [-0.9960098 ,  2.9960136 ],\n",
       "         [-0.99601364,  2.9960098 ]],\n",
       "\n",
       "        [[-0.9960098 ,  2.9960098 ],\n",
       "         [-0.9960098 ,  2.9960098 ],\n",
       "         [-0.9960098 ,  2.9960175 ]]]], dtype=float32)>"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "custom_bn_layer = custom_bn(scale=2, offset=1, axes=[3])\n",
    "x = custom_bn_layer(test_data, training = True)\n",
    "x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### *Testing: Dismantling the batch normalization process*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[[[ 0.5]\n",
      "   [ 2.5]\n",
      "   [ 4.5]]\n",
      "\n",
      "  [[ 6.5]\n",
      "   [ 8.5]\n",
      "   [10.5]]]\n",
      "\n",
      "\n",
      " [[[12.5]\n",
      "   [14.5]\n",
      "   [16.5]]\n",
      "\n",
      "  [[18.5]\n",
      "   [20.5]\n",
      "   [22.5]]]], shape=(2, 2, 3, 1), dtype=float32) tf.Tensor(\n",
      "[[[[0.25]\n",
      "   [0.25]\n",
      "   [0.25]]\n",
      "\n",
      "  [[0.25]\n",
      "   [0.25]\n",
      "   [0.25]]]\n",
      "\n",
      "\n",
      " [[[0.25]\n",
      "   [0.25]\n",
      "   [0.25]]\n",
      "\n",
      "  [[0.25]\n",
      "   [0.25]\n",
      "   [0.25]]]], shape=(2, 2, 3, 1), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# mean, var should have the shape [d1,d2,d3,1] where d1,d2,d3 are respective dimensions of the 4D input\n",
    "mean, variance = tf.nn.moments(test_data, axes = [3], keepdims=True)\n",
    "print(mean, variance)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2, 3, 2), dtype=float32, numpy=\n",
       "array([[[[-0.99601233,  2.9960122 ],\n",
       "         [-0.99601233,  2.9960122 ],\n",
       "         [-0.99601233,  2.9960122 ]],\n",
       "\n",
       "        [[-0.99601233,  2.9960122 ],\n",
       "         [-0.99601233,  2.9960122 ],\n",
       "         [-0.99601233,  2.9960122 ]]],\n",
       "\n",
       "\n",
       "       [[[-0.99601233,  2.9960122 ],\n",
       "         [-0.99601233,  2.9960122 ],\n",
       "         [-0.9960121 ,  2.9960122 ]],\n",
       "\n",
       "        [[-0.9960121 ,  2.9960122 ],\n",
       "         [-0.9960121 ,  2.9960122 ],\n",
       "         [-0.9960121 ,  2.9960122 ]]]], dtype=float32)>"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# using tf functions\n",
    "xhat = tf.add(tf.multiply(tf.divide(tf.subtract(test_data, mean), tf.math.sqrt(tf.add(variance, 0.001))), scale), offset)\n",
    "xhat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2, 3, 2), dtype=float32, numpy=\n",
       "array([[[[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]],\n",
       "\n",
       "        [[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]]],\n",
       "\n",
       "\n",
       "       [[[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]],\n",
       "\n",
       "        [[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]]]], dtype=float32)>"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.round(x-xhat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "    * Confirms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Similarly a list of axes can be passed.\n",
    "    * E.g. if axes = [0,2], first the moments are applied along 0, and then along 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2, 3, 2), dtype=float32, numpy=\n",
       "array([[[[-1.5730367 , -1.5730367 ],\n",
       "         [-0.9297775 , -0.92977744],\n",
       "         [-0.28651834, -0.28651834]],\n",
       "\n",
       "        [[-1.5730368 , -1.5730367 ],\n",
       "         [-0.9297776 , -0.9297776 ],\n",
       "         [-0.28651857, -0.28651834]]],\n",
       "\n",
       "\n",
       "       [[[ 2.2865183 ,  2.2865186 ],\n",
       "         [ 2.9297776 ,  2.9297776 ],\n",
       "         [ 3.5730367 ,  3.5730367 ]],\n",
       "\n",
       "        [[ 2.286518  ,  2.286518  ],\n",
       "         [ 2.9297771 ,  2.9297776 ],\n",
       "         [ 3.5730367 ,  3.5730367 ]]]], dtype=float32)>"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "custom_bn_layer = custom_bn(scale=2, offset=1, axes=[0,2])\n",
    "x = custom_bn_layer(test_data, training = True)\n",
    "x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### *Testing: Dismantling the batch normalization process*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[[[ 6.  7.]\n",
      "   [ 8.  9.]\n",
      "   [10. 11.]]\n",
      "\n",
      "  [[12. 13.]\n",
      "   [14. 15.]\n",
      "   [16. 17.]]]], shape=(1, 2, 3, 2), dtype=float32) tf.Tensor(\n",
      "[[[[36. 36.]\n",
      "   [36. 36.]\n",
      "   [36. 36.]]\n",
      "\n",
      "  [[36. 36.]\n",
      "   [36. 36.]\n",
      "   [36. 36.]]]], shape=(1, 2, 3, 2), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# Applying along 0 first resulting in shape [1,d2,d3,d4] \n",
    "mean, variance = tf.nn.moments(test_data, axes = [0], keepdims=True)\n",
    "print(mean, variance)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[[[ 8.  9.]]\n",
      "\n",
      "  [[14. 15.]]]], shape=(1, 2, 1, 2), dtype=float32) tf.Tensor(\n",
      "[[[[2.6666667 2.6666667]]\n",
      "\n",
      "  [[2.6666667 2.6666667]]]], shape=(1, 2, 1, 2), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# Applying along 2 to the mean from above \n",
    "mean_mean, variance_mean = tf.nn.moments(mean, axes = [2], keepdims=True)\n",
    "print(mean_mean, variance_mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[[[36. 36.]]\n",
      "\n",
      "  [[36. 36.]]]], shape=(1, 2, 1, 2), dtype=float32) tf.Tensor(\n",
      "[[[[0. 0.]]\n",
      "\n",
      "  [[0. 0.]]]], shape=(1, 2, 1, 2), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# Applying along 2 to the variance from above \n",
    "mean_var, var_var = tf.nn.moments(variance, axes = [2], keepdims=True)\n",
    "print(mean_var, var_var)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(1, 2, 1, 2), dtype=float32, numpy=\n",
       "array([[[[38.666668, 38.666668]],\n",
       "\n",
       "        [[38.666668, 38.666668]]]], dtype=float32)>"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# total variance is the variance from axes = 0 + var from axes = 2\n",
    "variance = mean_var + variance_mean\n",
    "variance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(1, 2, 1, 2), dtype=float32, numpy=\n",
       "array([[[[ 8.,  9.]],\n",
       "\n",
       "        [[14., 15.]]]], dtype=float32)>"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mean = mean_mean\n",
    "mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[[[ 8.  9.]]\n",
      "\n",
      "  [[14. 15.]]]], shape=(1, 2, 1, 2), dtype=float32) tf.Tensor(\n",
      "[[[[38.666668 38.666668]]\n",
      "\n",
      "  [[38.666668 38.666668]]]], shape=(1, 2, 1, 2), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# Applyling both together  \n",
    "mean, variance = tf.nn.moments(test_data, axes = [0,2], keepdims=True)\n",
    "print(mean, variance)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2, 3, 2), dtype=float32, numpy=\n",
       "array([[[[-1.5730367 , -1.5730367 ],\n",
       "         [-0.9297776 , -0.9297776 ],\n",
       "         [-0.28651834, -0.28651834]],\n",
       "\n",
       "        [[-1.5730367 , -1.5730367 ],\n",
       "         [-0.9297776 , -0.9297776 ],\n",
       "         [-0.28651834, -0.28651834]]],\n",
       "\n",
       "\n",
       "       [[[ 2.2865183 ,  2.2865183 ],\n",
       "         [ 2.9297776 ,  2.9297776 ],\n",
       "         [ 3.5730367 ,  3.5730367 ]],\n",
       "\n",
       "        [[ 2.2865183 ,  2.2865183 ],\n",
       "         [ 2.9297776 ,  2.9297776 ],\n",
       "         [ 3.5730367 ,  3.5730367 ]]]], dtype=float32)>"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# using tf functions\n",
    "xhat = tf.add(tf.multiply(tf.divide(tf.subtract(test_data, mean), tf.math.sqrt(tf.add(variance, 0.001))), scale), offset)\n",
    "xhat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2, 3, 2), dtype=float32, numpy=\n",
       "array([[[[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]],\n",
       "\n",
       "        [[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]]],\n",
       "\n",
       "\n",
       "       [[[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]],\n",
       "\n",
       "        [[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]]]], dtype=float32)>"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.round(x-xhat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "    * Confirms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Global normalization typically used for convolutional networks\n",
    "    * This is done by passing axes = [0,1,2] leaving out the last channel dimension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2, 3, 2), dtype=float32, numpy=\n",
       "array([[[[-2.1864765 , -2.1864765 ],\n",
       "         [-1.6071172 , -1.6071172 ],\n",
       "         [-1.0277578 , -1.0277579 ]],\n",
       "\n",
       "        [[-0.44839835, -0.4483986 ],\n",
       "         [ 0.13096094,  0.13096094],\n",
       "         [ 0.71032023,  0.71032023]]],\n",
       "\n",
       "\n",
       "       [[[ 1.2896798 ,  1.2896795 ],\n",
       "         [ 1.8690388 ,  1.869039  ],\n",
       "         [ 2.4483984 ,  2.448398  ]],\n",
       "\n",
       "        [[ 3.027758  ,  3.0277576 ],\n",
       "         [ 3.607117  ,  3.6071172 ],\n",
       "         [ 4.1864767 ,  4.186476  ]]]], dtype=float32)>"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "custom_bn_layer = custom_bn(scale=2, offset=1, axes=[0,1,2])\n",
    "x = custom_bn_layer(test_data, training = True)\n",
    "x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### *Testing: Dismantling the batch normalization process*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[[[ 6.  7.]\n",
      "   [ 8.  9.]\n",
      "   [10. 11.]]\n",
      "\n",
      "  [[12. 13.]\n",
      "   [14. 15.]\n",
      "   [16. 17.]]]], shape=(1, 2, 3, 2), dtype=float32) tf.Tensor(\n",
      "[[[[36. 36.]\n",
      "   [36. 36.]\n",
      "   [36. 36.]]\n",
      "\n",
      "  [[36. 36.]\n",
      "   [36. 36.]\n",
      "   [36. 36.]]]], shape=(1, 2, 3, 2), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# Applying along 0 first resulting in shape [1,d2,d3,d4] \n",
    "mean0, var0 = tf.nn.moments(test_data, axes = [0], keepdims=True)\n",
    "print(mean0, var0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[[[ 9. 10.]\n",
      "   [11. 12.]\n",
      "   [13. 14.]]]], shape=(1, 1, 3, 2), dtype=float32) tf.Tensor(\n",
      "[[[[9. 9.]\n",
      "   [9. 9.]\n",
      "   [9. 9.]]]], shape=(1, 1, 3, 2), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# Applying along 1 to mean0 from above \n",
    "mean1, var1 = tf.nn.moments(mean0, axes = [1], keepdims=True)\n",
    "print(mean1, var1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor([[[[11. 12.]]]], shape=(1, 1, 1, 2), dtype=float32) tf.Tensor([[[[2.6666667 2.6666667]]]], shape=(1, 1, 1, 2), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# Applying along 2 to the mean1 from above \n",
    "mean2, var2 = tf.nn.moments(mean1, axes = [2], keepdims=True)\n",
    "print(mean2, var2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(1, 1, 1, 2), dtype=float32, numpy=array([[[[11., 12.]]]], dtype=float32)>"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mean = mean2\n",
    "mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(1, 2, 3, 2), dtype=float32, numpy=\n",
       "array([[[[47.666668, 47.666668],\n",
       "         [47.666668, 47.666668],\n",
       "         [47.666668, 47.666668]],\n",
       "\n",
       "        [[47.666668, 47.666668],\n",
       "         [47.666668, 47.666668],\n",
       "         [47.666668, 47.666668]]]], dtype=float32)>"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# total variance is sum of all variances \n",
    "variance = var0+var1+var2\n",
    "variance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2,), dtype=float32, numpy=array([47.666668, 47.666668], dtype=float32)>"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "variance = tf.reduce_mean(variance, axis = [0,1,2])\n",
    "variance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor([[[[11. 12.]]]], shape=(1, 1, 1, 2), dtype=float32) tf.Tensor([[[[47.666668 47.666668]]]], shape=(1, 1, 1, 2), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# Applyling both together  \n",
    "mean, variance = tf.nn.moments(test_data, axes = [0,1,2], keepdims=True)\n",
    "print(mean, variance)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2, 3, 2), dtype=float32, numpy=\n",
       "array([[[[-2.1864765 , -2.1864765 ],\n",
       "         [-1.6071172 , -1.6071172 ],\n",
       "         [-1.0277579 , -1.0277579 ]],\n",
       "\n",
       "        [[-0.44839847, -0.44839847],\n",
       "         [ 0.13096094,  0.13096094],\n",
       "         [ 0.71032035,  0.71032035]]],\n",
       "\n",
       "\n",
       "       [[[ 1.2896796 ,  1.2896796 ],\n",
       "         [ 1.869039  ,  1.869039  ],\n",
       "         [ 2.4483986 ,  2.4483986 ]],\n",
       "\n",
       "        [[ 3.027758  ,  3.027758  ],\n",
       "         [ 3.6071172 ,  3.6071172 ],\n",
       "         [ 4.1864767 ,  4.1864767 ]]]], dtype=float32)>"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# using tf functions\n",
    "xhat = tf.add(tf.multiply(tf.divide(tf.subtract(test_data, mean), tf.math.sqrt(tf.add(variance,0.001))), scale), offset)\n",
    "xhat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2, 3, 2), dtype=float32, numpy=\n",
       "array([[[[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]],\n",
       "\n",
       "        [[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]]],\n",
       "\n",
       "\n",
       "       [[[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]],\n",
       "\n",
       "        [[0., 0.],\n",
       "         [0., 0.],\n",
       "         [0., 0.]]]], dtype=float32)>"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.round(x-xhat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "    * Confirms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Testing against Keras BatchNormalization Layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [],
   "source": [
    "def kerasbn(scale = 2, offset = 1, axis= [0]):\n",
    "    return tf.keras.layers.BatchNormalization(axis = axis, epsilon = 0.001, center = True, scale =True,\n",
    "                                             beta_initializer = tf.keras.initializers.constant(value = offset),\n",
    "                                             gamma_initializer = tf.keras.initializers.constant(value = scale))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [],
   "source": [
    "kerasbnlayer = kerasbn()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 3), dtype=float32, numpy=\n",
       "array([[ 2.9990008,  4.9980016,  6.9970026],\n",
       "       [ 8.996003 , 10.995004 , 12.994005 ]], dtype=float32)>"
      ]
     },
     "execution_count": 100,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kerasbnlayer(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 3), dtype=float32, numpy=\n",
       "array([[1., 2., 3.],\n",
       "       [4., 5., 6.]], dtype=float32)>"
      ]
     },
     "execution_count": 101,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_bn_layer = custom_bn(scale =2, offset = 1, axes=[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 3), dtype=float32, numpy=\n",
       "array([[-0.9995558, -0.9995556, -0.9995558],\n",
       "       [ 2.9995556,  2.999556 ,  2.9995556]], dtype=float32)>"
      ]
     },
     "execution_count": 103,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "custom_bn_layer(a, training =True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<tf.Variable 'batch_normalization_5/gamma:0' shape=(2,) dtype=float32, numpy=array([2., 2.], dtype=float32)>,\n",
       " <tf.Variable 'batch_normalization_5/beta:0' shape=(2,) dtype=float32, numpy=array([1., 1.], dtype=float32)>,\n",
       " <tf.Variable 'batch_normalization_5/moving_mean:0' shape=(2,) dtype=float32, numpy=array([0., 0.], dtype=float32)>,\n",
       " <tf.Variable 'batch_normalization_5/moving_variance:0' shape=(2,) dtype=float32, numpy=array([1., 1.], dtype=float32)>]"
      ]
     },
     "execution_count": 104,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kerasbnlayer.weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "    * Why is the output from the keras layer different?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 1), dtype=float32, numpy=\n",
       "array([[0.],\n",
       "       [0.]], dtype=float32)>"
      ]
     },
     "execution_count": 126,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kerasbnmean = kerasbnlayer.moving_mean\n",
    "kerasbnmean = tf.reshape(kerasbnmean, shape = (kerasbnmean.shape[0],1))\n",
    "kerasbnmean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "metadata": {},
   "outputs": [],
   "source": [
    "kerasbnvar = kerasbnlayer.moving_variance\n",
    "kerasbnvar = tf.reshape(kerasbnvar, shape = (kerasbnvar.shape[0],1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 3), dtype=float32, numpy=\n",
       "array([[ 2.9990008,  4.9980016,  6.9970026],\n",
       "       [ 8.996003 , 10.995004 , 12.994005 ]], dtype=float32)>"
      ]
     },
     "execution_count": 119,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.add(tf.multiply(tf.divide(tf.subtract(a, kerasbnmean), tf.math.sqrt(tf.add(kerasbnvar, 0.001))), 2), 1) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "    * The initialized moving mean and variance are being used in the calculation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<tf.Variable 'batch_normalization_5/gamma:0' shape=(2,) dtype=float32, numpy=array([2., 2.], dtype=float32)>,\n",
       " <tf.Variable 'batch_normalization_5/beta:0' shape=(2,) dtype=float32, numpy=array([1., 1.], dtype=float32)>]"
      ]
     },
     "execution_count": 120,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kerasbnlayer.trainable_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<tf.Variable 'batch_normalization_5/gamma:0' shape=(2,) dtype=float32, numpy=array([2., 2.], dtype=float32)>,\n",
       " <tf.Variable 'batch_normalization_5/beta:0' shape=(2,) dtype=float32, numpy=array([1., 1.], dtype=float32)>]"
      ]
     },
     "execution_count": 121,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kerasbnlayer.trainable_variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 122,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kerasbnlayer.trainable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (datascience)",
   "language": "python",
   "name": "datascience"
  },
  "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}