dropout4offnet.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": [
    "# tf.enable_eager_execution()"
   ]
  },
  {
   "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 = int(np.max(testing_data[1])) + 1\n",
    "number_of_input_channels = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "number_of_hidden_layers = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "number_of_candidate_layers = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "number_of_enabled_layers = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "number_of_features_per_layer = 8"
   ]
  },
  {
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "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": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dropout_offnet(images, number_of_candidate_layers, number_of_features_per_layer, weight_scale, number_of_enabled_layers, name=None):\n",
    "    with tf.variable_scope(name, default_name=\"dropout_offnet\"):\n",
    "        input_shape = images.shape\n",
    "        number_of_input_features = int(input_shape[3])\n",
    "        enabled_layer_ids = tf.random_shuffle(tf.range(number_of_candidate_layers))[:number_of_enabled_layers]\n",
    "        \n",
    "        output_indices = tf.concat(\n",
    "            (\n",
    "                tf.range(number_of_input_features),\n",
    "                tf.reshape(\n",
    "                    tf.expand_dims(enabled_layer_ids, axis=1) * number_of_features_per_layer +\n",
    "                        tf.expand_dims(tf.range(number_of_features_per_layer), axis=0),\n",
    "                    (number_of_features_per_layer * number_of_enabled_layers, )\n",
    "                ) + tf.constant(number_of_input_features, shape=(1, ))\n",
    "            ),\n",
    "            axis=0\n",
    "        )\n",
    "        number_of_features = int(number_of_features_per_layer * number_of_candidate_layers + number_of_input_features)\n",
    "        dense_weight_size = int((number_of_features - number_of_features_per_layer + number_of_input_features) * (number_of_features - number_of_input_features) // 2)\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, number_of_features_per_layer)\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=(number_of_candidate_layers, number_of_input_features, number_of_features_per_layer)\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=(number_of_candidate_layers, number_of_input_features, number_of_features_per_layer)\n",
    "        )\n",
    "        score_bias = tf.get_variable(\n",
    "            name=\"score_bias\",\n",
    "            initializer=tf.zeros_initializer(),\n",
    "            dtype=tf.float32,\n",
    "            shape=(number_of_candidate_layers, 1, 1, 1, number_of_features_per_layer),\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_candidate_layers, 1, 1, 1, number_of_features_per_layer),\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_candidate_layers, 1, 1, 1, number_of_features_per_layer),\n",
    "        )\n",
    "\n",
    "        def score_scale_initializer():\n",
    "            return tf.reshape(\n",
    "                tf.sqrt(weight_scale / tf.range(number_of_input_features, number_of_features, delta=number_of_features_per_layer, dtype=tf.float32)),\n",
    "                (number_of_candidate_layers, 1, 1)\n",
    "            )\n",
    "        score_scale = tf.get_variable(\n",
    "            name=\"score_scale\",\n",
    "            initializer=score_scale_initializer(),\n",
    "            dtype=tf.float32\n",
    "        )\n",
    "        def offset_scale_initializer():\n",
    "            return tf.fill(\n",
    "                (number_of_candidate_layers, 1, 1),\n",
    "                tf.sqrt(weight_scale / tf.cast(number_of_input_features, dtype=tf.float32))\n",
    "            )\n",
    "        offset_x_scale = tf.get_variable(\n",
    "            name=\"offset_x_scale\",\n",
    "            initializer=offset_scale_initializer(),\n",
    "            dtype=tf.float32\n",
    "        )\n",
    "        offset_y_scale = tf.get_variable(\n",
    "            name=\"offset_y_scale\",\n",
    "            initializer=offset_scale_initializer(),\n",
    "            dtype=tf.float32\n",
    "        )\n",
    "\n",
    "\n",
    "        def score_map(previous_features, i):\n",
    "            layer_id = enabled_layer_ids[i]\n",
    "            number_of_previous_features = number_of_input_features + layer_id * number_of_features_per_layer\n",
    "            weight_start = (number_of_input_features + number_of_previous_features) * (number_of_previous_features - number_of_input_features) // 2\n",
    "            tf.assert_rank(weight_start, 0)\n",
    "\n",
    "            # TODO test if the performance improves when using SparseTensor\n",
    "            tf.assert_rank(previous_features, 4)\n",
    "\n",
    "#             scale = tf.sqrt(weight_scale / tf.cast(number_of_input_features + i * number_of_features_per_layer, dtype=tf.float32))\n",
    "            \n",
    "            number_of_previous_enabled_features = number_of_input_features + i * number_of_features_per_layer\n",
    "\n",
    "            return grid_sample(\n",
    "                tf.cast(number_of_previous_features / number_of_previous_enabled_features, tf.float32) * (score_bias[layer_id] + tf.tensordot(\n",
    "                    previous_features,\n",
    "                    tf.gather(\n",
    "                        score_weight,\n",
    "                        weight_start + output_indices[:(number_of_input_features + i * number_of_features_per_layer)]\n",
    "                    ) * score_scale[layer_id],\n",
    "                    axes=1\n",
    "                )),\n",
    "                offset_x_bias[layer_id] + tf.tensordot(\n",
    "                    images,\n",
    "                    offset_x_weight[layer_id] * offset_x_scale[layer_id],\n",
    "                    axes=1\n",
    "                ),\n",
    "                offset_y_bias[layer_id] + tf.tensordot(\n",
    "                    images,\n",
    "                    offset_y_weight[layer_id] * offset_x_scale[layer_id],\n",
    "                    axes=1\n",
    "                ),\n",
    "            )\n",
    "\n",
    "        output, i = tf.while_loop(\n",
    "            lambda layers, i: i < number_of_enabled_layers,\n",
    "            lambda layers, i: (\n",
    "                tf.concat(\n",
    "                    (\n",
    "                        layers,\n",
    "                        score_map(\n",
    "                            layers,\n",
    "                            i\n",
    "                        )\n",
    "                    ),\n",
    "                    axis=3,\n",
    "                ),\n",
    "                i + 1\n",
    "            ),\n",
    "            (\n",
    "                images,\n",
    "                tf.constant(0, dtype=tf.int32),\n",
    "            ),\n",
    "            shape_invariants=(\n",
    "                tf.TensorShape((input_shape[0], input_shape[1], input_shape[2], None)), \n",
    "                tf.TensorShape(())\n",
    "            )\n",
    "        )\n",
    "        # tf.assert_equal(i, number_of_features)\n",
    "        \n",
    "        dynamic_input_shape = tf.shape(images)\n",
    "\n",
    "        return tf.transpose(tf.unsorted_segment_sum(tf.transpose(output), output_indices, number_of_input_features + number_of_features_per_layer * number_of_candidate_layers))\n",
    "#         return tf.IndexedSlices(\n",
    "#             output,\n",
    "#             output_indices,\n",
    "#             (dynamic_input_shape[0], dynamic_input_shape[1], dynamic_input_shape[2], number_of_input_features + number_of_features_per_layer * number_of_candidate_layers)\n",
    "#         )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def offnet(images, number_of_features, name=None):\n",
    "    with tf.variable_scope(name, default_name=\"offnet\"):\n",
    "        input_shape = images.shape\n",
    "#         if type(images) is tf.Tensor:\n",
    "#         elif type(images) is tf.IndexedSlices:\n",
    "#             input_shape = images.dense_shape\n",
    "#             images = tf.unsorted_segment_sum(\n",
    "#                 images.values,\n",
    "#                 images.indices,\n",
    "#                 input_shape[3],\n",
    "#             )\n",
    "            \n",
    "#             tf.IndexedSlices(\n",
    "#                 images.values,\n",
    "#                 images.indices,\n",
    "#                 tf.convert_to_tensor(images.dense_shape)\n",
    "#             ) # TODO: optimize this\n",
    "#         else:\n",
    "#             raise ValueError()\n",
    "        number_of_input_features = int(input_shape[3])    \n",
    "        score_weight = tf.get_variable(\n",
    "            name=\"score_weight\",\n",
    "            initializer=tf.random_normal_initializer(),\n",
    "            dtype=tf.float32,\n",
    "            shape=(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(),\n",
    "            dtype=tf.float32,\n",
    "            shape=(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(),\n",
    "            dtype=tf.float32,\n",
    "            shape=(number_of_input_features, number_of_features)\n",
    "        )\n",
    "        score_bias = tf.get_variable(\n",
    "            name=\"score_bias\",\n",
    "            initializer=tf.zeros_initializer(),\n",
    "            dtype=tf.float32,\n",
    "            shape=(1, 1, 1, number_of_features),\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=(1, 1, 1, number_of_features),\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=(1, 1, 1, number_of_features),\n",
    "        )\n",
    "        def scale_initializer():\n",
    "            return tf.reshape(tf.sqrt(weight_scale / number_of_input_features), (1, 1))\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",
    "        return grid_sample(\n",
    "            tf.tensordot(\n",
    "                images,\n",
    "                score_weight * score_scale,\n",
    "                axes=1\n",
    "            ) + score_bias,\n",
    "            tf.tensordot(\n",
    "                images,\n",
    "                offset_x_weight * offset_x_scale,\n",
    "                axes=1\n",
    "            ) + offset_x_bias,\n",
    "            tf.tensordot(\n",
    "                images,\n",
    "                offset_y_weight * offset_y_scale,\n",
    "                axes=1\n",
    "            ) + offset_y_bias,\n",
    "        )\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dense_offnet(images, number_of_additional_layers, number_of_features_per_layer, weight_scale, name=None):\n",
    "    with tf.variable_scope(name, default_name=\"dense_offnet\"):\n",
    "        input_shape = images.shape\n",
    "        number_of_input_features = int(input_shape[3])\n",
    "        number_of_features = int(number_of_features_per_layer * number_of_additional_layers + number_of_input_features)\n",
    "        dense_weight_size = int((number_of_features - number_of_features_per_layer + number_of_input_features) * (number_of_features - number_of_input_features) // 2)\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, number_of_features_per_layer)\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, number_of_features_per_layer)\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, number_of_features_per_layer)\n",
    "        )\n",
    "        score_bias = tf.get_variable(\n",
    "            name=\"score_bias\",\n",
    "            initializer=tf.zeros_initializer(),\n",
    "            dtype=tf.float32,\n",
    "            shape=(number_of_additional_layers, 1, 1, 1, number_of_features_per_layer),\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_additional_layers, 1, 1, 1, number_of_features_per_layer),\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_additional_layers, 1, 1, 1,  number_of_features_per_layer),\n",
    "        )\n",
    "\n",
    "        # TODO: index should be started at number_of_input_features\n",
    "        def scale_initializer():\n",
    "            return tf.sqrt(weight_scale / tf.range(number_of_input_features, number_of_features, delta=number_of_features_per_layer, dtype=tf.float32))\n",
    "\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",
    "\n",
    "        def score_map(previous_layers, i):\n",
    "            number_of_previous_layers = number_of_input_features + i * number_of_features_per_layer\n",
    "            weight_start = (number_of_input_features + number_of_previous_layers) * (number_of_previous_layers - number_of_input_features) // 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",
    "            tf.assert_rank(previous_layers, 4)\n",
    "\n",
    "            def indexed_sum(weight, bias, scale):\n",
    "                tf.assert_rank(weight, 2)\n",
    "                weight_slice = weight[weight_start:weight_end] * scale[i]\n",
    "\n",
    "                tf.assert_equal(tf.shape(weight_slice)[0], tf.shape(previous_layers)[3])\n",
    "                return tf.tensordot(\n",
    "                    previous_layers,\n",
    "                    weight_slice,\n",
    "                    axes=1\n",
    "                ) + bias[i]\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",
    "        output, i = tf.while_loop(\n",
    "            lambda layers, i: i < number_of_additional_layers,\n",
    "            lambda layers, i: (\n",
    "                tf.concat(\n",
    "                    (\n",
    "                        layers,\n",
    "                        score_map(\n",
    "                            layers,\n",
    "                            i\n",
    "                        )\n",
    "                    ),\n",
    "                    axis=3,\n",
    "                ),\n",
    "                i + 1\n",
    "            ),\n",
    "            (\n",
    "                images,\n",
    "                0,\n",
    "            ),\n",
    "            shape_invariants=(\n",
    "                tf.TensorShape((input_shape[0], input_shape[1], input_shape[2], None)), \n",
    "                tf.TensorShape(())\n",
    "            )\n",
    "        )\n",
    "        #tf.assert_equal(i, number_of_features)\n",
    "\n",
    "        \n",
    "        statically_shaped_output = output + 0 # Workaround for a tensorflow bug\n",
    "        statically_shaped_output.set_shape((input_shape[0], input_shape[1], input_shape[2], input_shape[3] + number_of_additional_layers * number_of_features_per_layer))\n",
    "        return statically_shaped_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "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",
    "    base_layers = dense_offnet(tf.expand_dims(padded_input, axis=3), number_of_hidden_layers, number_of_features_per_layer, weight_scale)\n",
    "    feature_layers = dropout_offnet(\n",
    "        base_layers,\n",
    "        number_of_candidate_layers,\n",
    "        number_of_features_per_layer,\n",
    "        weight_scale,\n",
    "        number_of_enabled_layers if mode == tf.estimator.ModeKeys.TRAIN else number_of_candidate_layers\n",
    "    )\n",
    "    layers = offnet(feature_layers, number_of_classes)\n",
    "\n",
    "    scores = tf.reduce_mean(layers, axis=(1, 2))\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",
    "    optimizer = tf.train.AdagradOptimizer(learning_rate=0.01)\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=tf.train.get_global_step()),\n",
    "            eval_metric_ops=eval_metric_ops,\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_config = tf.estimator.RunConfig(\n",
    "    model_dir=f\"models/dropout4offnet{number_of_features_per_layer}x{number_of_hidden_layers}+{number_of_features_per_layer}x{number_of_candidate_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": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using config: {'_model_dir': 'models/dropout4offnet8x32+8x100', '_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",
      ", '_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 0x7fa32bcb9908>, '_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 0x7fa32bcc5d08>) 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": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "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": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def testing_dataset():\n",
    "    return tf.data.Dataset.from_tensor_slices(testing_data).batch(batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "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 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/dropout4offnet8x32+8x100/model.ckpt.\n",
      "INFO:tensorflow:loss = 2.289009, step = 0\n",
      "INFO:tensorflow:global_step/sec: 0.884678\n",
      "INFO:tensorflow:loss = 2.1187534, step = 100 (113.037 sec)\n",
      "INFO:tensorflow:global_step/sec: 0.892884\n",
      "INFO:tensorflow:loss = 1.687089, step = 200 (111.996 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
}