bottleneckoffnet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mmh3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import functools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import iteration_utilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.python import debug as tf_debug"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from enum import IntEnum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "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": 11,
   "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": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "number_of_classes = int(np.max(testing_data[1])) + 1\n",
    "number_of_input_channels = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "number_of_experts = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "number_of_layers_per_block = 16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "number_of_features_per_layer = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "number_of_blocks = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "number_of_features_of_bottleneck = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "paddings = (0, 0), (2, 2), (2, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def offset_initializer():\n",
    "    return tf.random_normal_initializer(stddev=3.0)\n",
    "#     return tf.random_uniform_initializer(minval=-3.0, maxval=3.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def grid_sample(images, offset_x, offset_y):\n",
    "    tf.assert_rank(images, 4)\n",
    "    tf.assert_rank(offset_x, 4)\n",
    "    tf.assert_rank(offset_y, 4)\n",
    "    image_shape = tf.shape(images)\n",
    "    current_batch_size = image_shape[0]\n",
    "    width = image_shape[1]\n",
    "    height = image_shape[2]\n",
    "    number_of_features = image_shape[3]\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, feature_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",
    "        tf.range(number_of_features, dtype=tf.int32),\n",
    "        indexing='ij',\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_rank(left_index, 4)\n",
    "    tf.assert_rank(right_index, 4)\n",
    "    tf.assert_rank(bottom_index, 4)\n",
    "    tf.assert_rank(top_index, 4)\n",
    "    \n",
    "    images_top_left = tf.gather_nd(images, tf.stack((image_index, left_index, top_index, feature_index), axis=4))\n",
    "    images_top_right = tf.gather_nd(images, tf.stack((image_index, right_index, top_index, feature_index), axis=4))\n",
    "    images_bottom_left = tf.gather_nd(images, tf.stack((image_index, left_index, bottom_index, feature_index), axis=4))\n",
    "    images_bottom_right = tf.gather_nd(images, tf.stack((image_index, right_index, bottom_index, feature_index), axis=4))\n",
    "\n",
    "    tf.assert_rank(images_top_left, 4)\n",
    "    tf.assert_rank(images_top_right, 4)\n",
    "    tf.assert_rank(images_bottom_left, 4)\n",
    "    tf.assert_rank(images_bottom_right, 4)\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, 4)\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def offnet(images, number_of_features, number_of_experts=None, name=None):\n",
    "    with tf.variable_scope(name, default_name=\"offnet\"):\n",
    "        input_shape = images.shape\n",
    "        if number_of_experts is None:\n",
    "            number_of_experts = int(input_shape[0])\n",
    "        number_of_input_features = int(input_shape[4])\n",
    "        score_weight = tf.get_variable(\n",
    "            name=\"score_weight\",\n",
    "            initializer=tf.random_normal_initializer(stddev=tf.sqrt(1.0 / number_of_input_features)),\n",
    "            dtype=tf.float32,\n",
    "            shape=(number_of_experts, number_of_input_features, number_of_features)\n",
    "        )\n",
    "        offset_x_weight = tf.get_variable(\n",
    "            name=\"offset_x_weight\",\n",
    "            initializer=tf.random_normal_initializer(stddev=tf.sqrt(1.0 / number_of_input_features)),\n",
    "            dtype=tf.float32,\n",
    "            shape=(number_of_experts, number_of_input_features, number_of_features)\n",
    "        )\n",
    "        offset_y_weight = tf.get_variable(\n",
    "            name=\"offset_y_weight\",\n",
    "            initializer=tf.random_normal_initializer(stddev=tf.sqrt(1.0 / number_of_input_features)),\n",
    "            dtype=tf.float32,\n",
    "            shape=(number_of_experts, number_of_input_features, number_of_features)\n",
    "        )\n",
    "        offset_x_bias = tf.get_variable(\n",
    "            name=\"offset_x_bias\",\n",
    "            initializer=offset_initializer(),\n",
    "            dtype=tf.float32,\n",
    "            shape=(number_of_experts, 1, number_of_features),\n",
    "        )\n",
    "        offset_y_bias = tf.get_variable(\n",
    "            name=\"offset_y_bias\",\n",
    "            initializer=offset_initializer(),\n",
    "            dtype=tf.float32,\n",
    "            shape=(number_of_experts, 1,  number_of_features),\n",
    "        )\n",
    "#         def scale_initializer():\n",
    "#             return tf.sqrt(weight_scale / number_of_input_features)\n",
    "#         score_scale = tf.get_variable(\n",
    "#             name=\"score_scale\",\n",
    "#             initializer=scale_initializer(),\n",
    "#             dtype=tf.float32\n",
    "#         )\n",
    "#         offset_x_scale = tf.get_variable(\n",
    "#             name=\"offset_x_scale\",\n",
    "#             initializer=scale_initializer(),\n",
    "#             dtype=tf.float32\n",
    "#         )\n",
    "#         offset_y_scale = tf.get_variable(\n",
    "#             name=\"offset_y_scale\",\n",
    "#             initializer=scale_initializer(),\n",
    "#             dtype=tf.float32\n",
    "#         )\n",
    "        \n",
    "        dynamic_image_shape = tf.shape(images)\n",
    "        images_3d = tf.reshape(\n",
    "            images,\n",
    "            (\n",
    "                number_of_experts,\n",
    "                dynamic_image_shape[1] * dynamic_image_shape[2] * dynamic_image_shape[3],\n",
    "                number_of_input_features\n",
    "            )\n",
    "        )\n",
    "        def images_xw_plus_b(w, b):\n",
    "            tf.matmul(images_3d, w) + b\n",
    "        \n",
    "        def to_4d(images):\n",
    "            tf.assert_equal(tf.shape(images), (number_of_experts, dynamic_image_shape[1] * dynamic_image_shape[2] * dynamic_image_shape[3], number_of_features))\n",
    "            return tf.reshape(\n",
    "                images,\n",
    "                (\n",
    "                    number_of_experts * dynamic_image_shape[1],\n",
    "                    dynamic_image_shape[2],\n",
    "                    dynamic_image_shape[3],\n",
    "                    number_of_features\n",
    "                )\n",
    "            )\n",
    "        return tf.reshape(\n",
    "            grid_sample(\n",
    "                to_4d(tf.matmul(images_3d, score_weight)),\n",
    "                to_4d(tf.matmul(images_3d, offset_x_weight) + offset_x_bias),\n",
    "                to_4d(tf.matmul(images_3d, offset_y_weight) + offset_y_bias)\n",
    "            ),\n",
    "            (\n",
    "                number_of_experts,\n",
    "                dynamic_image_shape[1],\n",
    "                dynamic_image_shape[2],\n",
    "                dynamic_image_shape[3],\n",
    "                number_of_features\n",
    "            )\n",
    "        )        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def model_fn(features, labels, mode, params, config):\n",
    "    random.seed(number_of_features_per_layer * number_of_layers_per_block * number_of_blocks)\n",
    "    padded_input = tf.expand_dims(\n",
    "        tf.expand_dims(\n",
    "            tf.pad(tf.cast(features, tf.float32) / 255.0 - 0.5, paddings),\n",
    "            axis=0\n",
    "        ),\n",
    "        axis=4\n",
    "    )\n",
    "    \n",
    "    features = functools.reduce(\n",
    "        lambda images, i: offnet(\n",
    "            functools.reduce(\n",
    "                lambda features, i: tf.concat(\n",
    "                    (\n",
    "                        images,\n",
    "                        tf.layers.batch_normalization(\n",
    "                            offnet(images, number_of_features_per_layer, number_of_experts),\n",
    "                            training=mode is tf.estimator.ModeKeys.TRAIN\n",
    "                        )\n",
    "                    ),\n",
    "                    axis=4\n",
    "                ),\n",
    "                range(number_of_layers_per_block),\n",
    "                images,\n",
    "            ),\n",
    "            number_of_features_of_bottleneck,\n",
    "            number_of_experts\n",
    "        ),\n",
    "        range(number_of_blocks),\n",
    "        padded_input\n",
    "    )\n",
    "    \n",
    "    score_layer = offnet(features, number_of_classes, number_of_experts)\n",
    "\n",
    "    scores = tf.reduce_mean(\n",
    "        score_layer[:, :, 16, 16, :],\n",
    "        axis=0\n",
    "    )\n",
    "\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",
    "\n",
    "    optimizer = tf.train.AdamOptimizer(epsilon=0.1)\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",
    "        update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)\n",
    "        with tf.control_dependencies(update_ops):\n",
    "            return tf.estimator.EstimatorSpec(\n",
    "                mode=mode,\n",
    "                predictions=predictions,\n",
    "                loss=loss,\n",
    "                train_op=optimizer.minimize(loss, global_step=tf.train.get_global_step()),\n",
    "                eval_metric_ops=eval_metric_ops,\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_config = tf.estimator.RunConfig(\n",
    "    model_dir=f\"models/bottleneckoffnet{number_of_experts}x{number_of_features_per_layer}x{number_of_layers_per_block}x{number_of_blocks}\",\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": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using config: {'_model_dir': 'models/offnet1x8x5x3', '_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 0x7f81cf278f28>, '_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 0x7f81d2732a60>) 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": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "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": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def testing_dataset():\n",
    "    return tf.data.Dataset.from_tensor_slices(testing_data).batch(batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hook = tf_debug.TensorBoardDebugHook(\"localhost:6064\")\n",
    "# estimator.train(training_dataset,hooks=[hook])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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 1800 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/offnet1x8x5x3/model.ckpt.\n",
      "INFO:tensorflow:loss = 2.9007828, step = 0\n",
      "INFO:tensorflow:global_step/sec: 3.20163\n",
      "INFO:tensorflow:loss = 1.979384, step = 100 (31.234 sec)\n",
      "INFO:tensorflow:global_step/sec: 3.50959\n",
      "INFO:tensorflow:loss = 1.5797501, step = 200 (28.493 sec)\n",
      "INFO:tensorflow:global_step/sec: 3.5087\n",
      "INFO:tensorflow:loss = 1.2782334, step = 300 (28.501 sec)\n",
      "INFO:tensorflow:global_step/sec: 3.52153\n",
      "INFO:tensorflow:loss = 0.9823962, step = 400 (28.397 sec)\n",
      "INFO:tensorflow:global_step/sec: 3.51401\n",
      "INFO:tensorflow:loss = 0.9415303, step = 500 (28.457 sec)\n",
      "INFO:tensorflow:global_step/sec: 3.51597\n",
      "INFO:tensorflow:loss = 1.0150928, step = 600 (28.442 sec)\n",
      "INFO:tensorflow:global_step/sec: 3.52584\n",
      "INFO:tensorflow:loss = 0.88835955, step = 700 (28.362 sec)\n",
      "INFO:tensorflow:global_step/sec: 3.52371\n",
      "INFO:tensorflow:loss = 0.7851587, step = 800 (28.379 sec)\n",
      "INFO:tensorflow:global_step/sec: 3.52662\n",
      "INFO:tensorflow:loss = 0.7683656, step = 900 (28.356 sec)\n",
      "INFO:tensorflow:global_step/sec: 3.52499\n",
      "INFO:tensorflow:loss = 0.76362205, step = 1000 (28.369 sec)\n",
      "INFO:tensorflow:global_step/sec: 3.52027\n",
      "INFO:tensorflow:loss = 0.64977825, step = 1100 (28.407 sec)\n",
      "INFO:tensorflow:global_step/sec: 3.52465\n",
      "INFO:tensorflow:loss = 0.78920126, step = 1200 (28.371 sec)\n",
      "INFO:tensorflow:global_step/sec: 3.52072\n",
      "INFO:tensorflow:loss = 0.49506506, step = 1300 (28.403 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, throttle_secs=1800)\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
}