momentumoffnet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import abc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mmh3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import functools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import iteration_utilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.python import debug as tf_debug"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from enum import IntEnum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorFlow version: 1.8.0\n"
     ]
    }
   ],
   "source": [
    "print(\"TensorFlow version: {}\".format(tf.VERSION))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# (x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()\n",
    "training_data, testing_data = tf.keras.datasets.fashion_mnist.load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "number_of_classes = np.max(testing_data[1]).__int__() + 1\n",
    "number_of_input_channels = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "number_of_hidden_layers = 101"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "number_of_computing_layers = number_of_hidden_layers + number_of_classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "number_of_layers = number_of_input_channels + number_of_computing_layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "preferred_number_samples_per_layer = number_of_layers # All the previous layers "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "paddings = (0, 0), (2, 2), (2, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "random_seed = number_of_hidden_layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "weight_scale = 1.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6271\n"
     ]
    }
   ],
   "source": [
    "dense_weight_size = (number_of_layers + number_of_input_channels) * (number_of_layers - number_of_input_channels) // 2\n",
    "print(dense_weight_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def grid_sample(images, offset_x, offset_y):\n",
    "    tf.assert_rank(images, 3)\n",
    "    image_shape = tf.shape(images)\n",
    "    current_batch_size = image_shape[0]\n",
    "    width = image_shape[1]\n",
    "    height = image_shape[2]\n",
    "    \n",
    "    def check_width(tensor):\n",
    "        if (tf.test.is_gpu_available()):\n",
    "            return tensor\n",
    "        else:\n",
    "            return tf.maximum(tf.minimum(tensor, width - 1), 0)\n",
    "    \n",
    "    def check_height(tensor):\n",
    "        if (tf.test.is_gpu_available()):\n",
    "            return tensor\n",
    "        else:\n",
    "            return tf.maximum(tf.minimum(tensor, height - 1), 0)\n",
    "    \n",
    "    offset_left = tf.floor(offset_x)\n",
    "    offset_top = tf.floor(offset_y)\n",
    "    tf.assert_equal(tf.shape(offset_left), image_shape)\n",
    "    tf.assert_equal(tf.shape(offset_top), image_shape)\n",
    "\n",
    "    \n",
    "    eps = 1e-7\n",
    "\n",
    "    factor_right = offset_x - offset_left + eps\n",
    "    factor_left = 1.0 + 2.0 * eps - factor_right\n",
    "    factor_bottom = offset_y - offset_top + eps\n",
    "    factor_top = 1.0 + 2.0 * eps - factor_bottom\n",
    "    \n",
    "    image_index, x_index, y_index = tf.meshgrid(\n",
    "        tf.range(current_batch_size, dtype=tf.int32),\n",
    "        tf.range(width, dtype=tf.int32),\n",
    "        tf.range(height, dtype=tf.int32),\n",
    "        indexing='ij',\n",
    "    )\n",
    "    \n",
    "    tf.assert_equal(tf.shape(image_index), image_shape)\n",
    "    tf.assert_equal(tf.shape(x_index), image_shape)\n",
    "    tf.assert_equal(tf.shape(y_index), image_shape)\n",
    "\n",
    "    \n",
    "    left_index = check_width(x_index + tf.cast(offset_left, tf.int32))\n",
    "    top_index = check_height(y_index + tf.cast(offset_top, tf.int32))\n",
    "\n",
    "    right_index = left_index + 1\n",
    "    bottom_index = top_index + 1\n",
    "    \n",
    "    tf.assert_equal(tf.shape(left_index), image_shape)\n",
    "    tf.assert_equal(tf.shape(top_index), image_shape)\n",
    "    tf.assert_equal(tf.shape(right_index), image_shape)\n",
    "    tf.assert_equal(tf.shape(bottom_index), image_shape)    \n",
    "    \n",
    "    tf.assert_rank(left_index, 3)\n",
    "    tf.assert_rank(right_index, 3)\n",
    "    tf.assert_rank(bottom_index, 3)\n",
    "    tf.assert_rank(top_index, 3)\n",
    "    \n",
    "    images_top_left = tf.gather_nd(images, tf.stack((image_index, left_index, top_index), axis=3))\n",
    "    images_top_right = tf.gather_nd(images, tf.stack((image_index, right_index, top_index), axis=3))\n",
    "    images_bottom_left = tf.gather_nd(images, tf.stack((image_index, left_index, bottom_index), axis=3))\n",
    "    images_bottom_right = tf.gather_nd(images, tf.stack((image_index, right_index, bottom_index), axis=3))\n",
    "\n",
    "    tf.assert_rank(images_top_left, 3)\n",
    "    tf.assert_rank(images_top_right, 3)\n",
    "    tf.assert_rank(images_bottom_left, 3)\n",
    "    tf.assert_rank(images_bottom_right, 3)\n",
    "\n",
    "    lerp_top = factor_left * images_top_left + factor_right * images_top_right\n",
    "    lerp_bottom = factor_left * images_bottom_left + factor_right * images_bottom_right\n",
    "    output = factor_top * lerp_top + factor_bottom * lerp_bottom\n",
    "    tf.assert_rank(output, 3)\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def model_fn(features, labels, mode, params, config):\n",
    "    random.seed(number_of_hidden_layers)\n",
    "    padded_input = tf.pad(tf.cast(features, tf.float32) / 255.0 - 0.5, paddings)\n",
    "        \n",
    "    score_weight = tf.get_variable(\n",
    "        name=\"score_weight\",\n",
    "        initializer=tf.random_normal_initializer(),\n",
    "        dtype=tf.float32,\n",
    "        shape=(dense_weight_size, )\n",
    "    )\n",
    "    offset_x_weight = tf.get_variable(\n",
    "        name=\"offset_x_weight\",\n",
    "        initializer=tf.random_normal_initializer(),\n",
    "        dtype=tf.float32,\n",
    "        shape=(dense_weight_size, )\n",
    "    )\n",
    "    offset_y_weight = tf.get_variable(\n",
    "        name=\"offset_y_weight\",\n",
    "        initializer=tf.random_normal_initializer(),\n",
    "        dtype=tf.float32,\n",
    "        shape=(dense_weight_size, )\n",
    "    )\n",
    "    score_bias = tf.get_variable(\n",
    "        name=\"score_bias\",\n",
    "        initializer=tf.zeros_initializer(),\n",
    "        dtype=tf.float32,\n",
    "        shape=(number_of_layers, ),\n",
    "    )\n",
    "    offset_x_bias = tf.get_variable(\n",
    "        name=\"offset_x_bias\",\n",
    "        initializer=tf.random_uniform_initializer(minval=-3.0, maxval=3.0),\n",
    "        dtype=tf.float32,\n",
    "        shape=(number_of_layers, ),\n",
    "    )\n",
    "    offset_y_bias = tf.get_variable(\n",
    "        name=\"offset_y_bias\",\n",
    "        initializer=tf.random_uniform_initializer(minval=-3.0, maxval=3.0),\n",
    "        dtype=tf.float32,\n",
    "        shape=(number_of_layers, ),\n",
    "    )\n",
    "    score_scale = tf.get_variable(\n",
    "        name=\"score_scale\",\n",
    "        initializer=tf.sqrt(weight_scale / (tf.range(number_of_layers, dtype=tf.float32) + tf.constant(1e-10))),\n",
    "        dtype=tf.float32\n",
    "    )\n",
    "    offset_x_scale = tf.get_variable(\n",
    "        name=\"offset_x_scale\",\n",
    "        initializer=tf.sqrt(weight_scale / (tf.range(number_of_layers, dtype=tf.float32) + tf.constant(1e-10))),\n",
    "        dtype=tf.float32\n",
    "    )\n",
    "    offset_y_scale = tf.get_variable(\n",
    "        name=\"offset_y_scale\",\n",
    "        initializer=tf.sqrt(weight_scale / (tf.range(number_of_layers, dtype=tf.float32) + tf.constant(1e-10))),\n",
    "        dtype=tf.float32\n",
    "    )\n",
    "    \n",
    "\n",
    "    def score_map(previous_layers, number_of_previous_layers, next_layer_id):\n",
    "        weight_start = (number_of_input_channels + next_layer_id) * (next_layer_id - number_of_input_channels) // 2\n",
    "        weight_end = weight_start + number_of_previous_layers\n",
    "        tf.assert_rank(weight_start, 0)\n",
    "        \n",
    "        # TODO test if the performance improves when using SparseTensor\n",
    "        previous_layer_tensor = previous_layers.gather(tf.range(number_of_previous_layers))\n",
    "        tf.assert_rank(previous_layer_tensor, 4)\n",
    "\n",
    "        def indexed_sum(weight, bias, scale):\n",
    "            tf.assert_rank(weight, 1)\n",
    "            weight_slice = weight[weight_start:weight_end] * scale[next_layer_id]\n",
    "\n",
    "            tf.assert_equal(tf.shape(weight_slice)[0], tf.shape(previous_layer_tensor)[0])\n",
    "            return tf.tensordot(\n",
    "                weight_slice,\n",
    "                previous_layer_tensor,\n",
    "                axes=1\n",
    "            ) + bias[next_layer_id]\n",
    "\n",
    "        return grid_sample(\n",
    "            indexed_sum(score_weight, score_bias, score_scale),\n",
    "            indexed_sum(offset_x_weight, offset_x_bias, offset_x_scale),\n",
    "            indexed_sum(offset_y_weight, offset_y_bias, offset_y_scale),\n",
    "        )\n",
    "\n",
    "    layers, i = tf.while_loop(\n",
    "        lambda layers, i: i < number_of_layers,\n",
    "        lambda layers, i: (\n",
    "            layers.write(i,\n",
    "                score_map(\n",
    "                    layers,\n",
    "                    number_of_previous_layers=number_of_input_channels + number_of_hidden_layers,\n",
    "                    next_layer_id=i\n",
    "                )\n",
    "            ),\n",
    "            i + 1\n",
    "        ),\n",
    "        tf.while_loop(\n",
    "            lambda layers, i: i < number_of_input_channels + number_of_hidden_layers,\n",
    "            lambda layers, i: (layers.write(i, score_map(layers, i, i)), i + 1),\n",
    "            (\n",
    "                tf.TensorArray(tf.float32, size=number_of_layers, clear_after_read=False, infer_shape=True).write(\n",
    "                    0,\n",
    "                    padded_input,\n",
    "                ),\n",
    "                tf.constant(number_of_input_channels),\n",
    "            )\n",
    "        )\n",
    "    )\n",
    "\n",
    "    tf.assert_equal(i, tf.cast(number_of_layers, tf.int32))\n",
    "\n",
    "    scores_per_pixel = layers.gather(tf.range(number_of_layers - number_of_classes, number_of_layers, dtype=tf.int32))\n",
    "    score_shape = tf.shape(scores_per_pixel)\n",
    "    scores = tf.transpose(scores_per_pixel[:, :, score_shape[2] // 2, score_shape[3] // 2])\n",
    "    probabilities = tf.nn.softmax(logits=scores)\n",
    "    predicted_classes = tf.argmax(scores, 1)\n",
    "    predictions = {\n",
    "        'probabilities' : probabilities,\n",
    "        'scores': scores,\n",
    "        'class': predicted_classes,\n",
    "    }\n",
    "    global_step = tf.train.get_global_step()\n",
    "    optimizer = tf.train.MomentumOptimizer(\n",
    "        learning_rate=tf.train.exponential_decay(\n",
    "            0.01,\n",
    "            global_step=global_step,\n",
    "            decay_steps=10000,\n",
    "            decay_rate=0.95\n",
    "        ),\n",
    "        momentum=0.9,\n",
    "        use_nesterov=True\n",
    "    )\n",
    "    if labels is None:\n",
    "        return tf.estimator.EstimatorSpec(\n",
    "            mode=mode,\n",
    "            predictions=predictions,\n",
    "        )\n",
    "    else:\n",
    "        loss = tf.losses.softmax_cross_entropy(logits=scores, onehot_labels=tf.one_hot(labels, number_of_classes))\n",
    "        eval_metric_ops = {\n",
    "            'accuracy': tf.metrics.accuracy(labels=labels, predictions=predicted_classes)\n",
    "        }\n",
    "        return tf.estimator.EstimatorSpec(\n",
    "            mode=mode,\n",
    "            predictions=predictions,\n",
    "            loss=loss,\n",
    "            train_op=optimizer.minimize(loss, global_step=global_step),\n",
    "            eval_metric_ops=eval_metric_ops,\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_config = tf.estimator.RunConfig(\n",
    "    model_dir=f\"models/momentumoffnet{number_of_hidden_layers}\",\n",
    "    session_config=tf.ConfigProto(\n",
    "        gpu_options=tf.GPUOptions(\n",
    "            allow_growth=True,\n",
    "        ),\n",
    "        graph_options=tf.GraphOptions(\n",
    "            optimizer_options=tf.OptimizerOptions(\n",
    "                global_jit_level=tf.OptimizerOptions.ON_2,\n",
    "                do_function_inlining=True,\n",
    "                do_constant_folding=True,\n",
    "                do_common_subexpression_elimination=True,\n",
    "            ),\n",
    "        )\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using config: {'_model_dir': 'models/momentumoffnet101', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': gpu_options {\n",
      "  allow_growth: true\n",
      "}\n",
      "graph_options {\n",
      "  optimizer_options {\n",
      "    do_common_subexpression_elimination: true\n",
      "    do_constant_folding: true\n",
      "    do_function_inlining: true\n",
      "    global_jit_level: ON_2\n",
      "  }\n",
      "}\n",
      ", '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': None, '_service': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x7f6050e46160>, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1}\n",
      "WARNING:tensorflow:Estimator's model_fn (<function model_fn at 0x7f60e0043ea0>) includes params argument, but params are not passed to Estimator.\n"
     ]
    }
   ],
   "source": [
    "estimator = tf.estimator.Estimator(model_fn, config=run_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def training_dataset():\n",
    "    return tf.data.Dataset.from_tensor_slices(training_data).shuffle(1000).batch(batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def testing_dataset():\n",
    "    return tf.data.Dataset.from_tensor_slices(testing_data).batch(batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hook = tf_debug.TensorBoardDebugHook(\"localhost:6064\")\n",
    "# estimator.train(training_dataset,hooks=[hook])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Running training and evaluation locally (non-distributed).\n",
      "INFO:tensorflow:Start train and evaluate loop. The evaluate will happen after 600 secs (eval_spec.throttle_secs) or training is finished.\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Create CheckpointSaverHook.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "INFO:tensorflow:Saving checkpoints for 1 into models/momentumoffnet101/model.ckpt.\n",
      "INFO:tensorflow:loss = 2.2633867, step = 0\n",
      "INFO:tensorflow:global_step/sec: 2.2633\n",
      "INFO:tensorflow:loss = 1.1800361, step = 100 (44.184 sec)\n",
      "INFO:tensorflow:global_step/sec: 2.34612\n",
      "INFO:tensorflow:loss = 1.0103828, step = 200 (42.624 sec)\n",
      "INFO:tensorflow:global_step/sec: 2.32137\n",
      "INFO:tensorflow:loss = 0.9058804, step = 300 (43.078 sec)\n",
      "INFO:tensorflow:global_step/sec: 2.34554\n",
      "INFO:tensorflow:loss = 1.2011106, step = 400 (42.634 sec)\n",
      "INFO:tensorflow:global_step/sec: 2.31606\n",
      "INFO:tensorflow:loss = 2.5530915, step = 500 (43.177 sec)\n",
      "INFO:tensorflow:global_step/sec: 2.34621\n",
      "INFO:tensorflow:loss = 2.337885, step = 600 (42.622 sec)\n",
      "INFO:tensorflow:global_step/sec: 2.32692\n",
      "INFO:tensorflow:loss = 2.3227959, step = 700 (42.975 sec)\n",
      "INFO:tensorflow:global_step/sec: 2.276\n",
      "INFO:tensorflow:loss = 2.3114576, step = 800 (43.937 sec)\n",
      "INFO:tensorflow:global_step/sec: 2.32029\n",
      "INFO:tensorflow:loss = 2.2581215, step = 900 (43.098 sec)\n",
      "INFO:tensorflow:global_step/sec: 2.3226\n",
      "INFO:tensorflow:loss = 2.3056848, step = 1000 (43.055 sec)\n",
      "INFO:tensorflow:global_step/sec: 2.30674\n",
      "INFO:tensorflow:loss = 2.3734274, step = 1100 (43.352 sec)\n"
     ]
    }
   ],
   "source": [
    "tf.estimator.train_and_evaluate(\n",
    "    estimator,\n",
    "    train_spec=tf.estimator.TrainSpec(training_dataset),\n",
    "    eval_spec=tf.estimator.EvalSpec(testing_dataset)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# estimator.evaluate(input_fn=lambda:tf.data.Dataset.from_tensor_slices(testing_data).batch(3).take(1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# tuple(estimator.predict(input_fn=lambda:tf.data.Dataset.from_tensor_slices(testing_data).batch(3).take(1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# estimator.train(input_fn=training_dataset, steps=1)"
   ]
  },
  {
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}