notebooks/Tensorflow_MultiGPU.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#######################################################################################################\n",
    "# Summary\n",
    "# 1. Tensorflow Multi-GPU example using Estimator & Dataset high-APIs\n",
    "# 2. On-the-fly data-augmentation (random crop, random flip)\n",
    "# ToDo:\n",
    "# 3. Investigate tfrecord speed improvement (to match MXNet)\n",
    "# References:\n",
    "# https://www.tensorflow.org/performance/performance_guide\n",
    "# 1. https://jhui.github.io/2017/03/07/TensorFlow-Perforamnce-and-advance-topics/\n",
    "# 2. https://www.tensorflow.org/versions/master/performance/datasets_performance\n",
    "# 3. https://github.com/pudae/tensorflow-densenet\n",
    "# 4. https://stackoverflow.com/a/48096625/6772173\n",
    "# 5. https://stackoverflow.com/questions/47867748/transfer-learning-with-tf-estimator-estimator-framework\n",
    "# 6. https://github.com/BobLiu20/Classification_Nets/blob/master/tensorflow/common/average_gradients.py\n",
    "# 7. https://github.com/BobLiu20/Classification_Nets/blob/master/tensorflow/training/train_estimator.py\n",
    "#######################################################################################################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "MULTI_GPU = True  # TOGGLE THIS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash \n",
    "# Download model check-point and module from below repo:\n",
    "#wget -N https://github.com/pudae/tensorflow-densenet/raw/master/nets/densenet.py\n",
    "#wget -N https://ikpublictutorial.blob.core.windows.net/deeplearningframeworks/tf-densenet121.tar.gz\n",
    "#tar xzvf tf-densenet121.tar.gz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import time\n",
    "import multiprocessing\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from PIL import Image\n",
    "import random\n",
    "import tensorflow as tf\n",
    "from tensorflow.python.framework import dtypes\n",
    "from tensorflow.python.framework.ops import convert_to_tensor\n",
    "from tensorflow.contrib.data import Iterator\n",
    "from common.utils import download_data_chextxray, get_imgloc_labels, get_train_valid_test_split\n",
    "from common.utils import compute_roc_auc, get_cuda_version, get_cudnn_version, get_gpu_name\n",
    "from common.params_dense import *\n",
    "slim = tf.contrib.slim\n",
    "import densenet  # Download from https://github.com/pudae/tensorflow-densenet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OS:  linux\n",
      "Python:  3.5.2 |Anaconda custom (64-bit)| (default, Jul  2 2016, 17:53:06) \n",
      "[GCC 4.4.7 20120313 (Red Hat 4.4.7-1)]\n",
      "Numpy:  1.14.2\n",
      "Tensorflow:  1.6.0\n",
      "GPU:  ['Tesla P100-PCIE-16GB', 'Tesla P100-PCIE-16GB', 'Tesla P100-PCIE-16GB', 'Tesla P100-PCIE-16GB']\n",
      "CUDA Version 9.0.176\n",
      "CuDNN Version  7.0.5\n"
     ]
    }
   ],
   "source": [
    "print(\"OS: \", sys.platform)\n",
    "print(\"Python: \", sys.version)\n",
    "print(\"Numpy: \", np.__version__)\n",
    "print(\"Tensorflow: \", tf.__version__)\n",
    "print(\"GPU: \", get_gpu_name())\n",
    "print(get_cuda_version())\n",
    "print(\"CuDNN Version \", get_cudnn_version())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPUs:  24\n",
      "GPUs:  4\n"
     ]
    }
   ],
   "source": [
    "CPU_COUNT = multiprocessing.cpu_count()\n",
    "GPU_COUNT = len(get_gpu_name())\n",
    "print(\"CPUs: \", CPU_COUNT)\n",
    "print(\"GPUs: \", GPU_COUNT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "chestxray/images chestxray/Data_Entry_2017.csv\n"
     ]
    }
   ],
   "source": [
    "# Model-params\n",
    "IMAGENET_RGB_MEAN_CAFFE = np.array([123.68, 116.78, 103.94], dtype=np.float32)\n",
    "IMAGENET_SCALE_FACTOR_CAFFE = 0.017\n",
    "# Paths\n",
    "CSV_DEST = \"chestxray\"\n",
    "IMAGE_FOLDER = os.path.join(CSV_DEST, \"images\")\n",
    "LABEL_FILE = os.path.join(CSV_DEST, \"Data_Entry_2017.csv\")\n",
    "print(IMAGE_FOLDER, LABEL_FILE)\n",
    "CHKPOINT = 'tf-densenet121.ckpt'  # Downloaded tensorflow-checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Manually scale to multi-gpu\n",
    "if MULTI_GPU:\n",
    "    LR *= GPU_COUNT \n",
    "    BATCHSIZE *= GPU_COUNT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Please make sure to download\n",
      "https://docs.microsoft.com/en-us/azure/storage/common/storage-use-azcopy-linux#download-and-install-azcopy\n",
      "Data already exists\n",
      "CPU times: user 595 ms, sys: 289 ms, total: 885 ms\n",
      "Wall time: 883 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# Download data\n",
    "print(\"Please make sure to download\")\n",
    "print(\"https://docs.microsoft.com/en-us/azure/storage/common/storage-use-azcopy-linux#download-and-install-azcopy\")\n",
    "download_data_chextxray(CSV_DEST)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "#####################################################################################################\n",
    "## Data Loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class XrayData():\n",
    "    \n",
    "    def __init__(self, img_dir, lbl_file, patient_ids, mode, \n",
    "                 width=WIDTH, height=HEIGHT, batch_size=BATCHSIZE, \n",
    "                 imagenet_mean=IMAGENET_RGB_MEAN_CAFFE, imagenet_scaling = IMAGENET_SCALE_FACTOR_CAFFE,\n",
    "                 buffer=10):\n",
    "\n",
    "        self.img_locs, self.labels = get_imgloc_labels(img_dir, lbl_file, patient_ids)\n",
    "        self.data_size = len(self.labels)\n",
    "        self.imagenet_mean = imagenet_mean\n",
    "        self.imagenet_scaling = imagenet_scaling\n",
    "        self.width = width\n",
    "        self.height = height\n",
    "        data = tf.data.Dataset.from_tensor_slices((self.img_locs, self.labels))\n",
    "        \n",
    "        # Processing\n",
    "        # Output as channels-last and TF model will reshape in densenet.py\n",
    "        # inputs = tf.transpose(inputs, [0, 3, 1, 2])\n",
    "        if mode == 'training':\n",
    "            # Augmentation and repeat\n",
    "            data = data.shuffle(self.data_size).repeat().apply(\n",
    "                tf.contrib.data.map_and_batch(self._parse_function_train, batch_size)).prefetch(buffer)\n",
    "        elif mode == \"validation\":\n",
    "            # Repeat\n",
    "             data = data.repeat().apply(\n",
    "                tf.contrib.data.map_and_batch(self._parse_function_inference, batch_size)).prefetch(buffer)           \n",
    "        elif mode == 'testing':\n",
    "            # No repeat, no augmentation\n",
    "            data = data.apply(\n",
    "                tf.contrib.data.map_and_batch(self._parse_function_inference, batch_size)).prefetch(buffer)\n",
    "        \n",
    "        self.data = data        \n",
    "        print(\"Loaded {} labels and {} images\".format(len(self.labels), len(self.img_locs)))\n",
    "        \n",
    "        \n",
    "    def _parse_function_train(self, filename, label):\n",
    "        img_rgb, label = self._preprocess_image_labels(filename, label)\n",
    "        # Random crop (from 264x264)\n",
    "        img_rgb = tf.random_crop(img_rgb, [self.height, self.width, 3])\n",
    "        # Random flip\n",
    "        img_rgb = tf.image.random_flip_left_right(img_rgb)\n",
    "        # Channels-first\n",
    "        img_rgb = tf.transpose(img_rgb, [2, 0, 1])\n",
    "        return img_rgb, label\n",
    "        \n",
    "        \n",
    "    def _parse_function_inference(self, filename, label):\n",
    "        img_rgb, label = self._preprocess_image_labels(filename, label)\n",
    "        # Resize to final dimensions\n",
    "        img_rgb = tf.image.resize_images(img_rgb, [self.height, self.width])\n",
    "        # Channels-first\n",
    "        img_rgb = tf.transpose(img_rgb, [2, 0, 1])\n",
    "        return img_rgb, label \n",
    "       \n",
    "    \n",
    "    def _preprocess_image_labels(self, filename, label):\n",
    "        # load and preprocess the image\n",
    "        img_decoded = tf.to_float(tf.image.decode_png(tf.read_file(filename), channels=3))\n",
    "        img_centered = tf.subtract(img_decoded, self.imagenet_mean)\n",
    "        img_rgb = img_centered * self.imagenet_scaling\n",
    "        return img_rgb, tf.cast(label, dtype=tf.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train:21563 valid:3080 test:6162\n"
     ]
    }
   ],
   "source": [
    "train_set, valid_set, test_set = get_train_valid_test_split(TOT_PATIENT_NUMBER)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded 87306 labels and 87306 images\n",
      "Loaded 7616 labels and 7616 images\n",
      "Loaded 17198 labels and 17198 images\n"
     ]
    }
   ],
   "source": [
    "with tf.device('/cpu:0'):\n",
    "    # Create dataset for iterator\n",
    "    train_dataset = XrayData(img_dir=IMAGE_FOLDER, lbl_file=LABEL_FILE, patient_ids=train_set,  \n",
    "                             mode='training')\n",
    "    valid_dataset = XrayData(img_dir=IMAGE_FOLDER, lbl_file=LABEL_FILE, patient_ids=valid_set,\n",
    "                             mode='validation')\n",
    "    test_dataset  = XrayData(img_dir=IMAGE_FOLDER, lbl_file=LABEL_FILE, patient_ids=test_set,\n",
    "                             mode='testing')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "#####################################################################################################\n",
    "## Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def average_gradients(tower_grads):\n",
    "    average_grads = []\n",
    "    for grad_and_vars in zip(*tower_grads):\n",
    "        grads = []\n",
    "        for g, _ in grad_and_vars:\n",
    "            expanded_g = tf.expand_dims(g, 0)\n",
    "            grads.append(expanded_g)\n",
    "        grad = tf.concat(axis=0, values=grads)\n",
    "        grad = tf.reduce_mean(grad, 0)\n",
    "        v = grad_and_vars[0][1]\n",
    "        grad_and_var = (grad, v)\n",
    "        average_grads.append(grad_and_var)\n",
    "    return average_grads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_symbol(in_tensor, out_features):\n",
    "    # Import symbol\n",
    "    # is_training=True? | https://github.com/tensorflow/models/issues/3556\n",
    "    with slim.arg_scope(densenet.densenet_arg_scope(data_format=\"NCHW\")):\n",
    "        base_model, _ = densenet.densenet121(in_tensor,\n",
    "                                             num_classes=out_features,\n",
    "                                             is_training=True)\n",
    "        # Need to reshape from (?, 1, 1, 14) to (?, 14)\n",
    "        sym = tf.reshape(base_model, shape=[-1, out_features])\n",
    "    return sym"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_fn_single(features, labels, mode, params):\n",
    "    sym = get_symbol(features, out_features=params[\"n_classes\"])\n",
    "    # Predictions\n",
    "    predictions = tf.sigmoid(sym)\n",
    "    # ModeKeys.PREDICT\n",
    "    if mode == tf.estimator.ModeKeys.PREDICT:\n",
    "        return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)\n",
    "    # Optimizer & Loss\n",
    "    optimizer = tf.train.AdamOptimizer(params['lr'], beta1=0.9, beta2=0.999)\n",
    "    loss_fn = tf.losses.sigmoid_cross_entropy(labels, sym)\n",
    "    loss = tf.reduce_mean(loss_fn)\n",
    "    train_op = optimizer.minimize(loss, tf.train.get_global_step())\n",
    "    # Create eval metric ops\n",
    "    eval_metric_ops = {\"val_loss\": slim.metrics.streaming_mean(\n",
    "        tf.losses.sigmoid_cross_entropy(labels, predictions))}\n",
    "\n",
    "    return tf.estimator.EstimatorSpec(\n",
    "        mode=mode,\n",
    "        loss=loss,\n",
    "        train_op=train_op,\n",
    "        eval_metric_ops=eval_metric_ops)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def multi_gpu_X_y_split(features, labels, batchsize, gpus):\n",
    "    # Make sure splits sum to batch-size\n",
    "    split_size = batchsize // len(gpus)\n",
    "    splits = [split_size, ] * (len(gpus) - 1)\n",
    "    splits.append(batchsize - split_size * (len(gpus) - 1))\n",
    "    # Split the features and labels\n",
    "    features_split = tf.split(features, splits, axis=0)\n",
    "    labels_split = tf.split(labels, splits, axis=0)\n",
    "    return features_split, labels_split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_fn_multigpu(features, labels, mode, params):\n",
    "    if mode == tf.estimator.ModeKeys.PREDICT:\n",
    "        # Create symbol\n",
    "        sym = get_symbol(features, out_features=params[\"n_classes\"])\n",
    "        # Predictions\n",
    "        predictions = tf.sigmoid(sym)   \n",
    "        # ModeKeys.PREDICT\n",
    "        return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)\n",
    "    \n",
    "    # For multi-gpu split features and labels\n",
    "    features_split, labels_split = multi_gpu_X_y_split( features, labels, params[\"batchsize\"], params[\"gpus\"])\n",
    "    tower_grads = []\n",
    "    eval_logits = []\n",
    "    # Training operation\n",
    "    global_step = tf.train.get_global_step()\n",
    "    optimizer = tf.train.AdamOptimizer(LR, beta1=0.9, beta2=0.999)\n",
    "    # Load model on multiple GPUs\n",
    "    with tf.variable_scope(tf.get_variable_scope()):\n",
    "        for i in range(len(params['gpus'])):\n",
    "            with tf.device('/gpu:%d' % i), tf.name_scope('%s_%d' % (\"classification\", i)) as scope:\n",
    "                # Symbol\n",
    "                sym = get_symbol(features_split[i], out_features=params[\"n_classes\"])\n",
    "                # Loss\n",
    "                tf.losses.sigmoid_cross_entropy(labels_split[i], sym)\n",
    "                # Training-ops\n",
    "                update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS, scope)\n",
    "                updates_op = tf.group(*update_ops)\n",
    "                with tf.control_dependencies([updates_op]):\n",
    "                    losses = tf.get_collection(tf.GraphKeys.LOSSES, scope)\n",
    "                    total_loss = tf.add_n(losses, name='total_loss')\n",
    "                # reuse var\n",
    "                tf.get_variable_scope().reuse_variables()\n",
    "                # grad compute\n",
    "                grads = optimizer.compute_gradients(total_loss)\n",
    "                tower_grads.append(grads)\n",
    "                eval_logits.append(sym)\n",
    "\n",
    "    # We must calculate the mean of each gradient\n",
    "    grads = average_gradients(tower_grads)\n",
    "    # Apply the gradients to adjust the shared variables.\n",
    "    apply_gradient_op = optimizer.apply_gradients(grads, global_step=global_step)\n",
    "    # Group all updates to into a single train op.\n",
    "    train_op = tf.group(apply_gradient_op)\n",
    "    # Create eval metric ops (predict on multi-gpu)\n",
    "    predictions =  tf.concat(eval_logits, 0)\n",
    "    eval_metric_ops = {\"val_loss\": slim.metrics.streaming_mean(\n",
    "        tf.losses.sigmoid_cross_entropy(labels, predictions))}\n",
    "\n",
    "    return tf.estimator.EstimatorSpec(\n",
    "        mode=mode,\n",
    "        loss=total_loss,\n",
    "        train_op=train_op,\n",
    "        eval_metric_ops=eval_metric_ops)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_input_fn():\n",
    "    return train_dataset.data.make_one_shot_iterator().get_next()\n",
    "def valid_input_fn():\n",
    "    return valid_dataset.data.make_one_shot_iterator().get_next()\n",
    "def test_input_fn():\n",
    "    return test_dataset.data.make_one_shot_iterator().get_next()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Warm start from saved checkpoint (not logits)\n",
    "ws = tf.estimator.WarmStartSettings(ckpt_to_initialize_from=CHKPOINT, vars_to_warm_start=\"^(?!.*(logits))\")\n",
    "# Params\n",
    "params={\"lr\":LR, \"n_classes\":CLASSES, \"batchsize\":BATCHSIZE, \"gpus\":list(range(GPU_COUNT))}\n",
    "# Model functions\n",
    "if MULTI_GPU:\n",
    "    model_fn=model_fn_multigpu\n",
    "else:\n",
    "    model_fn=model_fn_single"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using default config.\n",
      "WARNING:tensorflow:Using temporary folder as model directory: /tmp/tmp4qzuvlp6\n",
      "INFO:tensorflow:Using config: {'_keep_checkpoint_max': 5, '_save_checkpoints_steps': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x7f328562a898>, '_task_id': 0, '_save_summary_steps': 100, '_task_type': 'worker', '_global_id_in_cluster': 0, '_session_config': None, '_is_chief': True, '_service': None, '_tf_random_seed': None, '_master': '', '_log_step_count_steps': 100, '_evaluation_master': '', '_num_ps_replicas': 0, '_keep_checkpoint_every_n_hours': 10000, '_model_dir': '/tmp/tmp4qzuvlp6', '_num_worker_replicas': 1, '_save_checkpoints_secs': 600}\n",
      "CPU times: user 1.35 ms, sys: 3.37 ms, total: 4.72 ms\n",
      "Wall time: 4.11 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# Create Estimator\n",
    "nn = tf.estimator.Estimator(model_fn=model_fn, params=params, warm_start_from=ws)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 31 µs, sys: 7 µs, total: 38 µs\n",
      "Wall time: 44.3 µs\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# Create train & eval specs\n",
    "train_spec = tf.estimator.TrainSpec(input_fn=train_input_fn,\n",
    "                                    max_steps=EPOCHS*(train_dataset.data_size//BATCHSIZE))\n",
    "# Hard to run validation every epoch so playing around with throttle_secs to get 5 runs\n",
    "eval_spec = tf.estimator.EvalSpec(input_fn=valid_input_fn,\n",
    "                                  throttle_secs=400)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "# 1 GPU - Main training loop: 50min 19s\n",
    "# 4 GPU - Main training loop: 25min 8s\n",
    "# Run train and evaluate (on validation data)\n",
    "tf.estimator.train_and_evaluate(nn, train_spec, eval_spec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Restoring parameters from /tmp/tmp4qzuvlp6/model.ckpt-1705\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "Full AUC [0.8095029096268271, 0.8730753531871049, 0.7890140063105355, 0.8841224641071695, 0.8791105523364449, 0.909663052586902, 0.7296254190733823, 0.846209840706696, 0.6351492058205271, 0.8433088569206826, 0.7720235409411627, 0.8095089838328521, 0.7474695735925051, 0.8898279027128161]\n",
      "Test AUC: 0.8155\n",
      "CPU times: user 3min 21s, sys: 21.9 s, total: 3min 43s\n",
      "Wall time: 44.4 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# Main prediction loop: 44.4s\n",
    "# Test AUC: 0.8155\n",
    "predictions = list(nn.predict(test_input_fn))\n",
    "y_truth = test_dataset.labels\n",
    "y_guess = np.array(predictions)\n",
    "print(\"Test AUC: {0:.4f}\".format(compute_roc_auc(y_truth, y_guess, CLASSES))) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "#####################################################################################################\n",
    "## Synthetic Data (Pure Training)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test on fake-data -> no IO lag\n",
    "batch_in_epoch = train_dataset.data_size//BATCHSIZE\n",
    "tot_num = batch_in_epoch * BATCHSIZE\n",
    "fake_X = np.random.rand(tot_num, 3, 224, 224).astype(np.float32)\n",
    "fake_y = np.random.rand(tot_num, CLASSES).astype(np.float32) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using default config.\n",
      "WARNING:tensorflow:Using temporary folder as model directory: /tmp/tmp81y2qplw\n",
      "INFO:tensorflow:Using config: {'_keep_checkpoint_max': 5, '_save_checkpoints_steps': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x7f327dd058d0>, '_task_id': 0, '_save_summary_steps': 100, '_task_type': 'worker', '_global_id_in_cluster': 0, '_session_config': None, '_is_chief': True, '_service': None, '_tf_random_seed': None, '_master': '', '_log_step_count_steps': 100, '_evaluation_master': '', '_num_ps_replicas': 0, '_keep_checkpoint_every_n_hours': 10000, '_model_dir': '/tmp/tmp81y2qplw', '_num_worker_replicas': 1, '_save_checkpoints_secs': 600}\n",
      "CPU times: user 5.27 ms, sys: 123 µs, total: 5.4 ms\n",
      "Wall time: 4.49 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# Create Estimator\n",
    "nn = tf.estimator.Estimator(model_fn=model_fn, params=params)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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 /tmp/tmp81y2qplw/model.ckpt.\n",
      "INFO:tensorflow:loss = 0.7307674, step = 1\n",
      "INFO:tensorflow:global_step/sec: 1.70092\n",
      "INFO:tensorflow:loss = 0.6955669, step = 101 (58.795 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.86103\n",
      "INFO:tensorflow:loss = 0.6935265, step = 201 (53.734 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.85679\n",
      "INFO:tensorflow:loss = 0.6932544, step = 301 (53.856 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.84197\n",
      "INFO:tensorflow:loss = 0.69336593, step = 401 (54.289 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.81324\n",
      "INFO:tensorflow:loss = 0.6918529, step = 501 (55.150 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.83523\n",
      "INFO:tensorflow:loss = 0.6934935, step = 601 (54.489 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.85034\n",
      "INFO:tensorflow:loss = 0.6935285, step = 701 (54.044 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.79776\n",
      "INFO:tensorflow:loss = 0.6927167, step = 801 (55.624 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.85051\n",
      "INFO:tensorflow:loss = 0.68868047, step = 901 (54.040 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.84869\n",
      "INFO:tensorflow:loss = 0.6937602, step = 1001 (54.091 sec)\n",
      "INFO:tensorflow:Saving checkpoints for 1088 into /tmp/tmp81y2qplw/model.ckpt.\n",
      "INFO:tensorflow:global_step/sec: 1.64897\n",
      "INFO:tensorflow:loss = 0.69516927, step = 1101 (60.643 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.80598\n",
      "INFO:tensorflow:loss = 0.69132197, step = 1201 (55.373 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.83902\n",
      "INFO:tensorflow:loss = 0.691399, step = 1301 (54.377 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.79873\n",
      "INFO:tensorflow:loss = 0.69149405, step = 1401 (55.595 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.81925\n",
      "INFO:tensorflow:loss = 0.6921871, step = 1501 (54.967 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.83108\n",
      "INFO:tensorflow:loss = 0.6896055, step = 1601 (54.613 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.7884\n",
      "INFO:tensorflow:loss = 0.69277227, step = 1701 (55.916 sec)\n",
      "INFO:tensorflow:Saving checkpoints for 1705 into /tmp/tmp81y2qplw/model.ckpt.\n",
      "INFO:tensorflow:Loss for final step: 0.69099534.\n",
      "CPU times: user 42min 44s, sys: 14min 11s, total: 56min 56s\n",
      "Wall time: 17min 6s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.estimator.estimator.Estimator at 0x7f3284ef0e80>"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "# 4 GPU - Synthetic data: 17min 6s\n",
    "nn.train(tf.estimator.inputs.numpy_input_fn(\n",
    "    fake_X,\n",
    "    fake_y,\n",
    "    shuffle=False,\n",
    "    num_epochs=EPOCHS,\n",
    "    batch_size=BATCHSIZE))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:py35]",
   "language": "python",
   "name": "conda-env-py35-py"
  },
  "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
}