7. train_custom_optimizer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### This tutorial demostrate how to design a custom optimizer.\n",
    "\n",
    "`Tim, Salimans, et al. \"Weight normalization: a simple reparameterization to accelerate training of deep neural networks\" 2016.`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "# import matplotlib for visualization\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "# set to plot in the notebook\n",
    "%matplotlib inline \n",
    "\n",
    "# import keras (using tensorflow backend)\n",
    "import keras"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from keras.backend.tensorflow_backend import set_session\n",
    "config = tf.ConfigProto()\n",
    "config.gpu_options.per_process_gpu_memory_fraction = 0.1\n",
    "config.gpu_options.visible_device_list= '0'\n",
    "set_session(tf.Session(config=config))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from keras.datasets import cifar10\n",
    "\n",
    "# set batch size and number of classes\n",
    "input_shape = (32,32,3)\n",
    "batch_size = 500\n",
    "num_classes = 10\n",
    "epochs = 10\n",
    "\n",
    "# load cifar10\n",
    "(x_train, y_train), (x_test, y_test) = cifar10.load_data()\n",
    "\n",
    "# convert class vectors to one-hot binary class matrices.\n",
    "y_train = keras.utils.to_categorical(y_train, num_classes)\n",
    "y_test = keras.utils.to_categorical(y_test, num_classes)\n",
    "\n",
    "# normalize data to [-1, 1]\n",
    "x_train = x_train.astype('float32')\n",
    "x_test = x_test.astype('float32')\n",
    "x_train /= 127.5\n",
    "x_test /= 127.5\n",
    "x_train -= 1.\n",
    "x_test -= 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "conv2d_1 (Conv2D)            (None, 30, 30, 32)        896       \n",
      "_________________________________________________________________\n",
      "activation_1 (Activation)    (None, 30, 30, 32)        0         \n",
      "_________________________________________________________________\n",
      "max_pooling2d_1 (MaxPooling2 (None, 15, 15, 32)        0         \n",
      "_________________________________________________________________\n",
      "activation_2 (Activation)    (None, 15, 15, 32)        0         \n",
      "_________________________________________________________________\n",
      "conv2d_2 (Conv2D)            (None, 13, 13, 64)        18496     \n",
      "_________________________________________________________________\n",
      "activation_3 (Activation)    (None, 13, 13, 64)        0         \n",
      "_________________________________________________________________\n",
      "max_pooling2d_2 (MaxPooling2 (None, 6, 6, 64)          0         \n",
      "_________________________________________________________________\n",
      "activation_4 (Activation)    (None, 6, 6, 64)          0         \n",
      "_________________________________________________________________\n",
      "flatten_1 (Flatten)          (None, 2304)              0         \n",
      "_________________________________________________________________\n",
      "dense_1 (Dense)              (None, 120)               276600    \n",
      "_________________________________________________________________\n",
      "activation_5 (Activation)    (None, 120)               0         \n",
      "_________________________________________________________________\n",
      "dense_2 (Dense)              (None, 10)                1210      \n",
      "_________________________________________________________________\n",
      "activation_6 (Activation)    (None, 10)                0         \n",
      "=================================================================\n",
      "Total params: 297,202\n",
      "Trainable params: 297,202\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "from keras.models import Sequential\n",
    "\n",
    "from keras.layers import Dense, Activation, Flatten, Conv2D, MaxPooling2D\n",
    "\n",
    "def cnn_model(input_shape, num_classes):    \n",
    "    model = Sequential()\n",
    "\n",
    "    # conv1\n",
    "    model.add(Conv2D(32, (3, 3), padding='valid', input_shape=input_shape))\n",
    "    model.add(Activation('relu'))\n",
    "    # pool1\n",
    "    model.add(MaxPooling2D(pool_size=(2,2)))\n",
    "    model.add(Activation('relu'))\n",
    "\n",
    "    # conv2\n",
    "    model.add(Conv2D(64, (3, 3), padding='valid'))\n",
    "    model.add(Activation('relu'))\n",
    "    # pool2\n",
    "    model.add(MaxPooling2D(pool_size=(2,2)))\n",
    "    model.add(Activation('relu'))\n",
    "\n",
    "    # fc1\n",
    "    model.add(Flatten())\n",
    "    model.add(Dense(120))\n",
    "    model.add(Activation('relu'))\n",
    "\n",
    "    # fc2\n",
    "    model.add(Dense(num_classes))\n",
    "    model.add(Activation('softmax'))\n",
    "    \n",
    "    return model\n",
    "\n",
    "model = cnn_model(input_shape, num_classes)\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<svg height=\"994pt\" viewBox=\"0.00 0.00 215.00 994.00\" width=\"215pt\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g class=\"graph\" id=\"graph0\" transform=\"scale(1 1) rotate(0) translate(4 990)\">\n",
       "<title>G</title>\n",
       "<polygon fill=\"white\" points=\"-4,4 -4,-990 211,-990 211,4 -4,4\" stroke=\"none\"/>\n",
       "<!-- 140196831456392 -->\n",
       "<g class=\"node\" id=\"node1\"><title>140196831456392</title>\n",
       "<polygon fill=\"none\" points=\"17.5,-949.5 17.5,-985.5 189.5,-985.5 189.5,-949.5 17.5,-949.5\" stroke=\"black\"/>\n",
       "<text font-family=\"Times,serif\" font-size=\"14.00\" text-anchor=\"middle\" x=\"103.5\" y=\"-963.8\">conv2d_1_input: InputLayer</text>\n",
       "</g>\n",
       "<!-- 140196831456280 -->\n",
       "<g class=\"node\" id=\"node2\"><title>140196831456280</title>\n",
       "<polygon fill=\"none\" points=\"42,-876.5 42,-912.5 165,-912.5 165,-876.5 42,-876.5\" stroke=\"black\"/>\n",
       "<text font-family=\"Times,serif\" font-size=\"14.00\" text-anchor=\"middle\" x=\"103.5\" y=\"-890.8\">conv2d_1: Conv2D</text>\n",
       "</g>\n",
       "<!-- 140196831456392&#45;&gt;140196831456280 -->\n",
       "<g class=\"edge\" id=\"edge1\"><title>140196831456392-&gt;140196831456280</title>\n",
       "<path d=\"M103.5,-949.313C103.5,-941.289 103.5,-931.547 103.5,-922.569\" fill=\"none\" stroke=\"black\"/>\n",
       "<polygon fill=\"black\" points=\"107,-922.529 103.5,-912.529 100,-922.529 107,-922.529\" stroke=\"black\"/>\n",
       "</g>\n",
       "<!-- 140196831456784 -->\n",
       "<g class=\"node\" id=\"node3\"><title>140196831456784</title>\n",
       "<polygon fill=\"none\" points=\"29.5,-803.5 29.5,-839.5 177.5,-839.5 177.5,-803.5 29.5,-803.5\" stroke=\"black\"/>\n",
       "<text font-family=\"Times,serif\" font-size=\"14.00\" text-anchor=\"middle\" x=\"103.5\" y=\"-817.8\">activation_1: Activation</text>\n",
       "</g>\n",
       "<!-- 140196831456280&#45;&gt;140196831456784 -->\n",
       "<g class=\"edge\" id=\"edge2\"><title>140196831456280-&gt;140196831456784</title>\n",
       "<path d=\"M103.5,-876.313C103.5,-868.289 103.5,-858.547 103.5,-849.569\" fill=\"none\" stroke=\"black\"/>\n",
       "<polygon fill=\"black\" points=\"107,-849.529 103.5,-839.529 100,-849.529 107,-849.529\" stroke=\"black\"/>\n",
       "</g>\n",
       "<!-- 140196831456896 -->\n",
       "<g class=\"node\" id=\"node4\"><title>140196831456896</title>\n",
       "<polygon fill=\"none\" points=\"0,-730.5 0,-766.5 207,-766.5 207,-730.5 0,-730.5\" stroke=\"black\"/>\n",
       "<text font-family=\"Times,serif\" font-size=\"14.00\" text-anchor=\"middle\" x=\"103.5\" y=\"-744.8\">max_pooling2d_1: MaxPooling2D</text>\n",
       "</g>\n",
       "<!-- 140196831456784&#45;&gt;140196831456896 -->\n",
       "<g class=\"edge\" id=\"edge3\"><title>140196831456784-&gt;140196831456896</title>\n",
       "<path d=\"M103.5,-803.313C103.5,-795.289 103.5,-785.547 103.5,-776.569\" fill=\"none\" stroke=\"black\"/>\n",
       "<polygon fill=\"black\" points=\"107,-776.529 103.5,-766.529 100,-776.529 107,-776.529\" stroke=\"black\"/>\n",
       "</g>\n",
       "<!-- 140196831462848 -->\n",
       "<g class=\"node\" id=\"node5\"><title>140196831462848</title>\n",
       "<polygon fill=\"none\" points=\"29.5,-657.5 29.5,-693.5 177.5,-693.5 177.5,-657.5 29.5,-657.5\" stroke=\"black\"/>\n",
       "<text font-family=\"Times,serif\" font-size=\"14.00\" text-anchor=\"middle\" x=\"103.5\" y=\"-671.8\">activation_2: Activation</text>\n",
       "</g>\n",
       "<!-- 140196831456896&#45;&gt;140196831462848 -->\n",
       "<g class=\"edge\" id=\"edge4\"><title>140196831456896-&gt;140196831462848</title>\n",
       "<path d=\"M103.5,-730.313C103.5,-722.289 103.5,-712.547 103.5,-703.569\" fill=\"none\" stroke=\"black\"/>\n",
       "<polygon fill=\"black\" points=\"107,-703.529 103.5,-693.529 100,-703.529 107,-703.529\" stroke=\"black\"/>\n",
       "</g>\n",
       "<!-- 140196831562216 -->\n",
       "<g class=\"node\" id=\"node6\"><title>140196831562216</title>\n",
       "<polygon fill=\"none\" points=\"42,-584.5 42,-620.5 165,-620.5 165,-584.5 42,-584.5\" stroke=\"black\"/>\n",
       "<text font-family=\"Times,serif\" font-size=\"14.00\" text-anchor=\"middle\" x=\"103.5\" y=\"-598.8\">conv2d_2: Conv2D</text>\n",
       "</g>\n",
       "<!-- 140196831462848&#45;&gt;140196831562216 -->\n",
       "<g class=\"edge\" id=\"edge5\"><title>140196831462848-&gt;140196831562216</title>\n",
       "<path d=\"M103.5,-657.313C103.5,-649.289 103.5,-639.547 103.5,-630.569\" fill=\"none\" stroke=\"black\"/>\n",
       "<polygon fill=\"black\" points=\"107,-630.529 103.5,-620.529 100,-630.529 107,-630.529\" stroke=\"black\"/>\n",
       "</g>\n",
       "<!-- 140196831563672 -->\n",
       "<g class=\"node\" id=\"node7\"><title>140196831563672</title>\n",
       "<polygon fill=\"none\" points=\"29.5,-511.5 29.5,-547.5 177.5,-547.5 177.5,-511.5 29.5,-511.5\" stroke=\"black\"/>\n",
       "<text font-family=\"Times,serif\" font-size=\"14.00\" text-anchor=\"middle\" x=\"103.5\" y=\"-525.8\">activation_3: Activation</text>\n",
       "</g>\n",
       "<!-- 140196831562216&#45;&gt;140196831563672 -->\n",
       "<g class=\"edge\" id=\"edge6\"><title>140196831562216-&gt;140196831563672</title>\n",
       "<path d=\"M103.5,-584.313C103.5,-576.289 103.5,-566.547 103.5,-557.569\" fill=\"none\" stroke=\"black\"/>\n",
       "<polygon fill=\"black\" points=\"107,-557.529 103.5,-547.529 100,-557.529 107,-557.529\" stroke=\"black\"/>\n",
       "</g>\n",
       "<!-- 140196797551112 -->\n",
       "<g class=\"node\" id=\"node8\"><title>140196797551112</title>\n",
       "<polygon fill=\"none\" points=\"0,-438.5 0,-474.5 207,-474.5 207,-438.5 0,-438.5\" stroke=\"black\"/>\n",
       "<text font-family=\"Times,serif\" font-size=\"14.00\" text-anchor=\"middle\" x=\"103.5\" y=\"-452.8\">max_pooling2d_2: MaxPooling2D</text>\n",
       "</g>\n",
       "<!-- 140196831563672&#45;&gt;140196797551112 -->\n",
       "<g class=\"edge\" id=\"edge7\"><title>140196831563672-&gt;140196797551112</title>\n",
       "<path d=\"M103.5,-511.313C103.5,-503.289 103.5,-493.547 103.5,-484.569\" fill=\"none\" stroke=\"black\"/>\n",
       "<polygon fill=\"black\" points=\"107,-484.529 103.5,-474.529 100,-484.529 107,-484.529\" stroke=\"black\"/>\n",
       "</g>\n",
       "<!-- 140196797552848 -->\n",
       "<g class=\"node\" id=\"node9\"><title>140196797552848</title>\n",
       "<polygon fill=\"none\" points=\"29.5,-365.5 29.5,-401.5 177.5,-401.5 177.5,-365.5 29.5,-365.5\" stroke=\"black\"/>\n",
       "<text font-family=\"Times,serif\" font-size=\"14.00\" text-anchor=\"middle\" x=\"103.5\" y=\"-379.8\">activation_4: Activation</text>\n",
       "</g>\n",
       "<!-- 140196797551112&#45;&gt;140196797552848 -->\n",
       "<g class=\"edge\" id=\"edge8\"><title>140196797551112-&gt;140196797552848</title>\n",
       "<path d=\"M103.5,-438.313C103.5,-430.289 103.5,-420.547 103.5,-411.569\" fill=\"none\" stroke=\"black\"/>\n",
       "<polygon fill=\"black\" points=\"107,-411.529 103.5,-401.529 100,-411.529 107,-411.529\" stroke=\"black\"/>\n",
       "</g>\n",
       "<!-- 140196797619784 -->\n",
       "<g class=\"node\" id=\"node10\"><title>140196797619784</title>\n",
       "<polygon fill=\"none\" points=\"48.5,-292.5 48.5,-328.5 158.5,-328.5 158.5,-292.5 48.5,-292.5\" stroke=\"black\"/>\n",
       "<text font-family=\"Times,serif\" font-size=\"14.00\" text-anchor=\"middle\" x=\"103.5\" y=\"-306.8\">flatten_1: Flatten</text>\n",
       "</g>\n",
       "<!-- 140196797552848&#45;&gt;140196797619784 -->\n",
       "<g class=\"edge\" id=\"edge9\"><title>140196797552848-&gt;140196797619784</title>\n",
       "<path d=\"M103.5,-365.313C103.5,-357.289 103.5,-347.547 103.5,-338.569\" fill=\"none\" stroke=\"black\"/>\n",
       "<polygon fill=\"black\" points=\"107,-338.529 103.5,-328.529 100,-338.529 107,-338.529\" stroke=\"black\"/>\n",
       "</g>\n",
       "<!-- 140196797619448 -->\n",
       "<g class=\"node\" id=\"node11\"><title>140196797619448</title>\n",
       "<polygon fill=\"none\" points=\"52.5,-219.5 52.5,-255.5 154.5,-255.5 154.5,-219.5 52.5,-219.5\" stroke=\"black\"/>\n",
       "<text font-family=\"Times,serif\" font-size=\"14.00\" text-anchor=\"middle\" x=\"103.5\" y=\"-233.8\">dense_1: Dense</text>\n",
       "</g>\n",
       "<!-- 140196797619784&#45;&gt;140196797619448 -->\n",
       "<g class=\"edge\" id=\"edge10\"><title>140196797619784-&gt;140196797619448</title>\n",
       "<path d=\"M103.5,-292.313C103.5,-284.289 103.5,-274.547 103.5,-265.569\" fill=\"none\" stroke=\"black\"/>\n",
       "<polygon fill=\"black\" points=\"107,-265.529 103.5,-255.529 100,-265.529 107,-265.529\" stroke=\"black\"/>\n",
       "</g>\n",
       "<!-- 140196797685040 -->\n",
       "<g class=\"node\" id=\"node12\"><title>140196797685040</title>\n",
       "<polygon fill=\"none\" points=\"29.5,-146.5 29.5,-182.5 177.5,-182.5 177.5,-146.5 29.5,-146.5\" stroke=\"black\"/>\n",
       "<text font-family=\"Times,serif\" font-size=\"14.00\" text-anchor=\"middle\" x=\"103.5\" y=\"-160.8\">activation_5: Activation</text>\n",
       "</g>\n",
       "<!-- 140196797619448&#45;&gt;140196797685040 -->\n",
       "<g class=\"edge\" id=\"edge11\"><title>140196797619448-&gt;140196797685040</title>\n",
       "<path d=\"M103.5,-219.313C103.5,-211.289 103.5,-201.547 103.5,-192.569\" fill=\"none\" stroke=\"black\"/>\n",
       "<polygon fill=\"black\" points=\"107,-192.529 103.5,-182.529 100,-192.529 107,-192.529\" stroke=\"black\"/>\n",
       "</g>\n",
       "<!-- 140196797684816 -->\n",
       "<g class=\"node\" id=\"node13\"><title>140196797684816</title>\n",
       "<polygon fill=\"none\" points=\"52.5,-73.5 52.5,-109.5 154.5,-109.5 154.5,-73.5 52.5,-73.5\" stroke=\"black\"/>\n",
       "<text font-family=\"Times,serif\" font-size=\"14.00\" text-anchor=\"middle\" x=\"103.5\" y=\"-87.8\">dense_2: Dense</text>\n",
       "</g>\n",
       "<!-- 140196797685040&#45;&gt;140196797684816 -->\n",
       "<g class=\"edge\" id=\"edge12\"><title>140196797685040-&gt;140196797684816</title>\n",
       "<path d=\"M103.5,-146.313C103.5,-138.289 103.5,-128.547 103.5,-119.569\" fill=\"none\" stroke=\"black\"/>\n",
       "<polygon fill=\"black\" points=\"107,-119.529 103.5,-109.529 100,-119.529 107,-119.529\" stroke=\"black\"/>\n",
       "</g>\n",
       "<!-- 140196797323920 -->\n",
       "<g class=\"node\" id=\"node14\"><title>140196797323920</title>\n",
       "<polygon fill=\"none\" points=\"29.5,-0.5 29.5,-36.5 177.5,-36.5 177.5,-0.5 29.5,-0.5\" stroke=\"black\"/>\n",
       "<text font-family=\"Times,serif\" font-size=\"14.00\" text-anchor=\"middle\" x=\"103.5\" y=\"-14.8\">activation_6: Activation</text>\n",
       "</g>\n",
       "<!-- 140196797684816&#45;&gt;140196797323920 -->\n",
       "<g class=\"edge\" id=\"edge13\"><title>140196797684816-&gt;140196797323920</title>\n",
       "<path d=\"M103.5,-73.3129C103.5,-65.2895 103.5,-55.5475 103.5,-46.5691\" fill=\"none\" stroke=\"black\"/>\n",
       "<polygon fill=\"black\" points=\"107,-46.5288 103.5,-36.5288 100,-46.5289 107,-46.5288\" stroke=\"black\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>"
      ],
      "text/plain": [
       "<IPython.core.display.SVG object>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# show model\n",
    "from IPython.display import SVG\n",
    "from keras.utils.vis_utils import model_to_dot\n",
    "\n",
    "SVG(model_to_dot(model).create(prog='dot', format='svg'))\n",
    "\n",
    "# or you can save model diagram by following code\n",
    "#from keras.utils import plot_model\n",
    "#plot_model(model, to_file='cnn_model.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10\n",
      "100/100 [==============================] - 3s 26ms/step - loss: 1.6251 - categorical_accuracy: 0.4227 - val_loss: 1.3572 - val_categorical_accuracy: 0.5211\n",
      "Epoch 2/10\n",
      "100/100 [==============================] - 2s 17ms/step - loss: 1.2700 - categorical_accuracy: 0.5544 - val_loss: 1.2177 - val_categorical_accuracy: 0.5656\n",
      "Epoch 3/10\n",
      "100/100 [==============================] - 2s 16ms/step - loss: 1.1356 - categorical_accuracy: 0.6043 - val_loss: 1.1058 - val_categorical_accuracy: 0.6116\n",
      "Epoch 4/10\n",
      "100/100 [==============================] - 2s 15ms/step - loss: 1.0451 - categorical_accuracy: 0.6362 - val_loss: 1.0827 - val_categorical_accuracy: 0.6206\n",
      "Epoch 5/10\n",
      "100/100 [==============================] - 2s 15ms/step - loss: 0.9816 - categorical_accuracy: 0.6600 - val_loss: 1.0072 - val_categorical_accuracy: 0.6492\n",
      "Epoch 6/10\n",
      "100/100 [==============================] - 2s 15ms/step - loss: 0.9246 - categorical_accuracy: 0.6812 - val_loss: 0.9668 - val_categorical_accuracy: 0.6658\n",
      "Epoch 7/10\n",
      "100/100 [==============================] - 2s 16ms/step - loss: 0.8819 - categorical_accuracy: 0.6956 - val_loss: 0.9400 - val_categorical_accuracy: 0.6776\n",
      "Epoch 8/10\n",
      "100/100 [==============================] - 2s 15ms/step - loss: 0.8400 - categorical_accuracy: 0.7126 - val_loss: 0.9252 - val_categorical_accuracy: 0.6863\n",
      "Epoch 9/10\n",
      "100/100 [==============================] - 2s 15ms/step - loss: 0.8176 - categorical_accuracy: 0.7183 - val_loss: 0.9176 - val_categorical_accuracy: 0.6873\n",
      "Epoch 10/10\n",
      "100/100 [==============================] - 2s 16ms/step - loss: 0.7779 - categorical_accuracy: 0.7337 - val_loss: 0.9253 - val_categorical_accuracy: 0.6829\n"
     ]
    }
   ],
   "source": [
    "from keras import backend as K\n",
    "from keras.optimizers import Adam\n",
    "\n",
    "# adapted from keras.optimizers.Adam\n",
    "class custom_optimizer(Adam):\n",
    "    def get_updates(self, params, loss):\n",
    "        grads = self.get_gradients(loss, params)\n",
    "        self.updates = [K.update_add(self.iterations, 1)]\n",
    "\n",
    "        lr = self.lr\n",
    "        if self.initial_decay > 0:\n",
    "            lr *= (1. / (1. + self.decay * tf.cast(self.iterations, tf.float32)))\n",
    "\n",
    "        t = tf.cast(self.iterations, tf.float32) + 1\n",
    "        lr_t = lr * K.sqrt(1. - K.pow(self.beta_2, t)) / (1. - K.pow(self.beta_1, t))\n",
    "\n",
    "        shapes = [K.get_variable_shape(p) for p in params]\n",
    "        ms = [K.zeros(shape) for shape in shapes]\n",
    "        vs = [K.zeros(shape) for shape in shapes]\n",
    "        self.weights = [self.iterations] + ms + vs\n",
    "\n",
    "        for p, g, m, v in zip(params, grads, ms, vs):\n",
    "\n",
    "            # if a weight tensor (len > 1) use weight normalized parameterization\n",
    "            ps = K.get_variable_shape(p)\n",
    "            if len(ps)>1:\n",
    "\n",
    "                # get weight normalization parameters\n",
    "                V, V_norm, V_scaler, g_param, grad_g, grad_V = get_custom_params_and_grads(p, g)\n",
    "\n",
    "                # Adam containers for the 'g' parameter\n",
    "                V_scaler_shape = K.get_variable_shape(V_scaler)\n",
    "                m_g = K.zeros(V_scaler_shape)\n",
    "                v_g = K.zeros(V_scaler_shape)\n",
    "\n",
    "                # update g parameters\n",
    "                m_g_t = (self.beta_1 * m_g) + (1. - self.beta_1) * grad_g\n",
    "                v_g_t = (self.beta_2 * v_g) + (1. - self.beta_2) * K.square(grad_g)\n",
    "                new_g_param = g_param - lr_t * m_g_t / (K.sqrt(v_g_t) + self.epsilon)\n",
    "                self.updates.append(K.update(m_g, m_g_t))\n",
    "                self.updates.append(K.update(v_g, v_g_t))\n",
    "\n",
    "                # update V parameters\n",
    "                m_t = (self.beta_1 * m) + (1. - self.beta_1) * grad_V\n",
    "                v_t = (self.beta_2 * v) + (1. - self.beta_2) * K.square(grad_V)\n",
    "                new_V_param = V - lr_t * m_t / (K.sqrt(v_t) + self.epsilon)\n",
    "                self.updates.append(K.update(m, m_t))\n",
    "                self.updates.append(K.update(v, v_t))\n",
    "\n",
    "                # wn param updates --> W updates\n",
    "                add_custom_param_updates(self.updates, new_V_param, new_g_param, p, V_scaler)\n",
    "\n",
    "            else: # do optimization normally\n",
    "                m_t = (self.beta_1 * m) + (1. - self.beta_1) * g\n",
    "                v_t = (self.beta_2 * v) + (1. - self.beta_2) * K.square(g)\n",
    "                p_t = p - lr_t * m_t / (K.sqrt(v_t) + self.epsilon)\n",
    "\n",
    "                self.updates.append(K.update(m, m_t))\n",
    "                self.updates.append(K.update(v, v_t))\n",
    "\n",
    "                new_p = p_t\n",
    "                #new_p = K.clip(new_p,-1.,1.)\n",
    "                \n",
    "                self.updates.append(K.update(p, new_p))\n",
    "        return self.updates\n",
    "    \n",
    "def get_custom_params_and_grads(p, g):\n",
    "    ps = K.get_variable_shape(p)\n",
    "\n",
    "    # construct weight scaler: V_scaler = g/||V||\n",
    "    V_scaler_shape = (ps[-1],)  # assumes we're using tensorflow!\n",
    "    V_scaler = K.ones(V_scaler_shape)  # init to ones, so effective parameters don't change\n",
    "\n",
    "    # get V parameters = ||V||/g * W\n",
    "    norm_axes = [i for i in range(len(ps) - 1)]\n",
    "    V = p / tf.reshape(V_scaler, [1] * len(norm_axes) + [-1])\n",
    "\n",
    "    # split V_scaler into ||V|| and g parameters\n",
    "    V_norm = tf.sqrt(tf.reduce_sum(tf.square(V), norm_axes))\n",
    "    g_param = V_scaler * V_norm\n",
    "\n",
    "    # get grad in V,g parameters\n",
    "    grad_g = tf.reduce_sum(g * V, norm_axes) / V_norm\n",
    "    grad_V = tf.reshape(V_scaler, [1] * len(norm_axes) + [-1]) * \\\n",
    "             (g - tf.reshape(grad_g / V_norm, [1] * len(norm_axes) + [-1]) * V)\n",
    "\n",
    "    return V, V_norm, V_scaler, g_param, grad_g, grad_V\n",
    "\n",
    "\n",
    "def add_custom_param_updates(updates, new_V_param, new_g_param, W, V_scaler):\n",
    "    ps = K.get_variable_shape(new_V_param)\n",
    "    norm_axes = [i for i in range(len(ps) - 1)]\n",
    "\n",
    "    # update W and V_scaler\n",
    "    new_V_norm = tf.sqrt(tf.reduce_sum(tf.square(new_V_param), norm_axes))\n",
    "    new_V_scaler = new_g_param / new_V_norm\n",
    "    new_W = tf.reshape(new_V_scaler, [1] * len(norm_axes) + [-1]) * new_V_param\n",
    "    #new_W = K.clip(new_W,-1.,1.)\n",
    "    \n",
    "    updates.append(K.update(W, new_W))\n",
    "    updates.append(K.update(V_scaler, new_V_scaler))\n",
    "    \n",
    "optimizer = custom_optimizer()\n",
    "\n",
    "# compile model\n",
    "loss = keras.losses.categorical_crossentropy\n",
    "metric = keras.metrics.categorical_accuracy\n",
    "model.compile(loss=loss,optimizer=optimizer,metrics=[metric])\n",
    "\n",
    "# import image data generator\n",
    "from keras.preprocessing.image import ImageDataGenerator\n",
    "datagen = ImageDataGenerator()\n",
    "\n",
    "# Compute quantities required for feature-wise normalization\n",
    "datagen.fit(x_train)\n",
    "\n",
    "# Fit the model on the batches generated by datagen.flow().\n",
    "history = model.fit_generator(\n",
    "    datagen.flow(x_train, y_train, batch_size=batch_size),\n",
    "    steps_per_epoch=x_train.shape[0] // batch_size,\n",
    "    epochs=epochs,\n",
    "    validation_data=(x_test, y_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train accuracy:73.89%.\n"
     ]
    }
   ],
   "source": [
    "[train_loss, train_acc] = history.model.evaluate(x=x_train, y=y_train, batch_size=batch_size, verbose=0)\n",
    "print(\"Train accuracy:{0:.2f}%.\".format(train_acc*100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test accuracy:68.29%.\n"
     ]
    }
   ],
   "source": [
    "[test_loss, test_acc] = history.model.evaluate(x=x_test, y=y_test, batch_size=batch_size, verbose=0)\n",
    "print(\"Test accuracy:{0:.2f}%.\".format(test_acc*100))"
   ]
  },
  {
   "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}