train_caenet.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "train_caenet.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/starkdg/pyConvnetPhash/blob/master/train_caenet.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "metadata": {
        "id": "jjcK55t-TFMq",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/gdrive')\n",
        "\n",
        "import sys\n",
        "import os\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.image as mpimg\n",
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "\n",
        "model_dir = \"/gdrive/My Drive/models\"\n",
        "module_inception_url = \"https://tfhub.dev/google/imagenet/inception_v3/feature_vector/1\"\n",
        "module_mobilenetv2_url = \"https://tfhub.dev/google/imagenet/mobilenet_v2_140_224/feature_vector/2\"\n",
        "\n",
        "module = hub.Module(module_mobilenetv2_url)\n",
        "target_height, target_width = hub.get_expected_image_size(module)\n",
        "\n",
        "n_inputs = 1792\n",
        "n_hidden = 256\n",
        "\n",
        "training_files_dir = \"/gdrive/My Drive/imageset/train\"\n",
        "validation_files_dir = \"/gdrive/My Drive/imageset/validation\"\n",
        "testing_files_dir = \"/gdrive/My Drive/imageset/test\"\n",
        "\n",
        "tied_weights=True\n",
        "\n",
        "batch_size = 10\n",
        "epochs = 10\n",
        "steps = 2000\n",
        "learning_rate = 0.001\n",
        "lambda_reg = 0.001\n",
        "model_tag = 1\n",
        "frozen_model = \"/gdrive/My Drive/models/cae_autoencoder/mobilenetv2_cae_autoenc_{0}to{1}_frozen_model-{2}.pb\".format(n_inputs, n_hidden, model_tag)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Ih_PGfUWLF95",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def get_tfrecord_files(path):\n",
        "  files = []\n",
        "  for entry in os.scandir(path):\n",
        "    if entry.is_file() and entry.name.endswith('.tfrecord'):\n",
        "             files.append(entry.path)\n",
        "  return files\n",
        "  \n",
        "def _parse_example(example):\n",
        "  features = {'height': tf.FixedLenFeature([], tf.int64),\n",
        "              'width': tf.FixedLenFeature([], tf.int64),\n",
        "              'image_raw': tf.FixedLenFeature([], tf.string)}\n",
        "  parsed_features = tf.parse_single_example(example, features)\n",
        "  img = tf.io.decode_raw(parsed_features['image_raw'], tf.uint8)\n",
        "  height = tf.cast(parsed_features['height'], tf.int32)\n",
        "  width = tf.cast(parsed_features['width'], tf.int32)\n",
        "\n",
        "  img_reshaped = tf.manip.reshape(img, [height, width, 3])\n",
        "  imgfl = tf.image.convert_image_dtype(img_reshaped, dtype=tf.float32)\n",
        "  img_norm = tf.expand_dims(imgfl, 0)\n",
        "  img_resized = tf.image.resize_bicubic(img_norm, [target_height, target_width])\n",
        "  img_resized = tf.squeeze(img_resized, 0)\n",
        "  return img_resized\n",
        "\n",
        "def input_function(path, batch_size=1, num_epochs=None, shuffle=False):\n",
        "  tfrecords = get_tfrecord_files(path)\n",
        "  dataset = tf.data.TFRecordDataset(tfrecords)\n",
        "  dataset = dataset.map(_parse_example)\n",
        "  if (shuffle):\n",
        "    dataset = dataset.shuffle(10000)\n",
        "  dataset = dataset.batch(batch_size).repeat(num_epochs)\n",
        "  iterator = dataset.make_initializable_iterator()\n",
        "  return iterator"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "OzKxh7rdTjgJ",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def get_weights():\n",
        "  weights = dict()\n",
        "  with tf.variable_scope(\"weights\", reuse=tf.AUTO_REUSE):\n",
        "    weights['w1'] = tf.get_variable('w1', shape=[n_inputs, n_hidden], trainable=True, initializer=tf.contrib.layers.xavier_initializer())\n",
        "    weights['b1'] = tf.Variable(tf.zeros([n_hidden], dtype=tf.float32))\n",
        "    if tied_weights:\n",
        "      weights['w2'] = tf.transpose(weights['w1'])\n",
        "    else:\n",
        "      weights['w2'] = tf.get_variable('w2', shape=[n_hidden, n_inputs], trainable=True, initializer=tf.contrib.layers.xavier_initializer())\n",
        "    weights['b2'] = tf.Variable(tf.zeros([n_inputs], dtype=tf.float32)) \n",
        "  return weights"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "xd_MYyIBcgcy",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "def create_cae_autoencoder(learning_rate, lambda_reg):\n",
        "  \n",
        "  weights = get_weights()\n",
        "  \n",
        "  x = tf.placeholder(tf.float32, shape=(None, n_inputs), name=\"input\")  \n",
        "  \n",
        "  num_x = tf.subtract(x, tf.reduce_min(x))\n",
        "  den_x = tf.subtract(tf.reduce_max(x), tf.reduce_min(x))\n",
        "  x_norm = tf.math.xdivy(num_x, den_x, name=\"normalization\")\n",
        "  \n",
        "  reg_term = tf.constant(lambda_reg, dtype=tf.float32, name=\"jnorm_scale_factor\")\n",
        "  \n",
        "  # input_dims -> n_hidden\n",
        "  h = tf.nn.sigmoid(tf.add(tf.matmul(x_norm, weights['w1']), weights['b1']), name=\"output256\")\n",
        "    \n",
        "  # reconstruction  n_hidden ->  input_dims\n",
        "  y = tf.identity(tf.add(tf.matmul(h, weights['w2']), weights['b2']), name=\"reconstructed\")\n",
        "   \n",
        "  # Jacobian norm\n",
        "  dhi = tf.square(tf.multiply(h, tf.subtract(tf.constant(1.0, tf.float32), h))) # N x n_hidden\n",
        "  dwj = tf.reduce_sum(tf.square(weights['w2']), axis=1, keepdims=True)          # n_hidden x n_input => n_hidden x 1\n",
        "  jnorm = tf.matmul(dhi, dwj, name=\"jnorm\")                                     # N x 1 \n",
        "  jnorm_mean = tf.reduce_mean(jnorm)\n",
        "  \n",
        "  # cost \n",
        "  cost = tf.reduce_sum(tf.nn.sigmoid_cross_entropy_with_logits(labels=x_norm, logits=y), axis=1, keepdims=True)\n",
        "  avg_cost = tf.reduce_mean(tf.add(cost, tf.multiply(reg_term, jnorm)))\n",
        "\n",
        "  with tf.variable_scope('opt', reuse=tf.AUTO_REUSE):\n",
        "    optimizer = tf.train.AdamOptimizer(learning_rate) \n",
        "    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)\n",
        "    with tf.control_dependencies(update_ops):\n",
        "      opt = optimizer.minimize(avg_cost)\n",
        "     \n",
        "  return x, h, weights, avg_cost, jnorm_mean, opt\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "p8BcbOitLQ5l",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def train_cae_model(training_files_dir,\n",
        "                        validation_files_dir,\n",
        "                        testing_files_dir,\n",
        "                        batch_size, epochs, steps, learning_rate, lambda_reg):\n",
        "    period_size = 100\n",
        "    input_dims = n_inputs\n",
        "    \n",
        "    training_iter = input_function(training_files_dir, batch_size)\n",
        "    training_images = training_iter.get_next()\n",
        "    training_features = module(training_images)\n",
        "        \n",
        "    validation_iter = input_function(validation_files_dir, batch_size)\n",
        "    validation_images = validation_iter.get_next()\n",
        "    validation_features = module(validation_images)\n",
        "        \n",
        "    testing_iter = input_function(testing_files_dir, 100)\n",
        "    testing_images = testing_iter.get_next()\n",
        "    testing_features = module(testing_images)\n",
        "        \n",
        "    x, out1, weights, recon_cost, jnorm, train_op = create_cae_autoencoder(learning_rate, lambda_reg)\n",
        "   \n",
        "    init = tf.global_variables_initializer()\n",
        "    sess = tf.Session()\n",
        "    sess.run(init)\n",
        "  \n",
        "    train_losses = []\n",
        "    valid_losses = []\n",
        "    valid_jlosses = []\n",
        "    print(\"Train CAE Autoencoder model\")\n",
        "    for i in range(epochs):\n",
        "      sess.run([training_iter.initializer, validation_iter.initializer])\n",
        "      iteration = 0\n",
        "      total_cost = 0.\n",
        "      total_val_cost = 0.\n",
        "      total_train_jcost = 0.\n",
        "      total_valid_jcost = 0.\n",
        "      while True:\n",
        "        try:\n",
        "          Xtrain = sess.run(training_features)\n",
        "          train_cost, _ = sess.run([recon_cost, train_op], feed_dict={x: Xtrain})\n",
        "          if (iteration % period_size == 0):\n",
        "            Xvalid = sess.run(validation_features)\n",
        "            validation_cost, valid_jcost = sess.run([recon_cost, jnorm], feed_dict={x: Xvalid})\n",
        "            total_cost += train_cost\n",
        "            total_val_cost += validation_cost\n",
        "            total_valid_jcost += valid_jcost\n",
        "          iteration = iteration + 1\n",
        "        except tf.errors.OutOfRangeError:\n",
        "          break\n",
        "        if (iteration > steps):\n",
        "          break\n",
        "                \n",
        "      steps_taken = iteration//period_size\n",
        "      avg_train_loss = total_cost/steps_taken\n",
        "      avg_val_loss = total_val_cost/steps_taken\n",
        "      avg_valid_jcost = total_valid_jcost/steps_taken\n",
        "      print(\"epoch {0} training cost {1:.4f} valid. cost {2:.4f} (jnorm {3})\".format(i+1, avg_train_loss, avg_val_loss, avg_valid_jcost))\n",
        "      train_losses.append(avg_train_loss)\n",
        "      valid_losses.append(avg_val_loss)\n",
        "      valid_jlosses.append(avg_valid_jcost)\n",
        "                         \n",
        "    plt.figure(101)\n",
        "    plt.plot(train_losses)\n",
        "    plt.plot(valid_losses)\n",
        "    plt.plot(valid_jlosses)\n",
        "    title_str = \"CAE Autoencoder {0} -> {1}\".format(n_inputs, n_hidden)\n",
        "    plt.title(title_str)\n",
        "    plt.xlabel(\"epochs\")\n",
        "    plt.ylabel(\"cost\")\n",
        "    plt.legend([\"training\", \"validation\", \"jnorm\"], loc=\"upper right\")\n",
        "    plt.show()\n",
        "    \n",
        "    w1 = weights['w1'].eval(session=sess)\n",
        "    w2 = weights['w2'].eval(session=sess)\n",
        "    b1 = weights['b1'].eval(session=sess)\n",
        "    b2 = weights['b2'].eval(session=sess)\n",
        "    \n",
        "    plt.figure(102)\n",
        "    plt.hist((w1.ravel(), w2.ravel()), bins=100, histtype='bar', stacked=True)\n",
        "    plt.legend(['w1', 'w2'], loc='upper right')\n",
        "    plt.title('Histogram of Weights')\n",
        "    plt.show()\n",
        "    \n",
        "    plt.figure(103)\n",
        "    plt.hist((b1.ravel(), b2.ravel()), bins=100, histtype='bar', stacked=True)\n",
        "    plt.legend(['b1', 'b2'], loc='upper right')\n",
        "    plt.title('Histogram of Biases')\n",
        "    plt.show()\n",
        "    \n",
        "    print(\"run test on 100 images\")\n",
        "    sess.run([testing_iter.initializer])\n",
        "    Xtest = sess.run(testing_features)\n",
        "    testing_cost, testing_jnorm = sess.run([recon_cost, jnorm], feed_dict={x: Xtest})\n",
        "    print(\"test cost = {0:.4f} jnorm = {1}\".format(testing_cost, testing_jnorm))\n",
        "    \n",
        "    graphdef = tf.get_default_graph().as_graph_def()\n",
        "    subgraphdef = tf.graph_util.extract_sub_graph(graphdef, ['output256'])\n",
        "    subgraphdef = tf.graph_util.remove_training_nodes(subgraphdef)\n",
        "    subgraphdef_frozen = tf.graph_util.convert_variables_to_constants(sess, subgraphdef, ['output256'])\n",
        "    \n",
        "    print(\"write model: \", frozen_model)\n",
        "    with tf.gfile.GFile(frozen_model, \"wb\") as f:\n",
        "      f.write(subgraphdef_frozen.SerializeToString())\n",
        "        \n",
        "    sess.close() "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "ycizNOxjVVDb",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "print(\"Train autoencoder\")\n",
        "print(\"training files: \", training_files_dir)\n",
        "print(\"validation files: \", validation_files_dir)\n",
        "print(\"testing files: \", testing_files_dir)\n",
        "print(\"learning_rate: \", learning_rate)\n",
        "print(\"batch size: \", batch_size)\n",
        "print(\"epochs: \", epochs)\n",
        "print(\"steps: \", steps)\n",
        "print(\"lambda reg: \", lambda_reg)\n",
        "\n",
        "train_cae_model(training_files_dir,\n",
        "                    validation_files_dir,\n",
        "                    testing_files_dir,\n",
        "                    batch_size, epochs, steps,\n",
        "                    learning_rate, lambda_reg)\n"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}