mnist-magnet-loss.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting MNIST_data/train-images-idx3-ubyte.gz\n",
      "Extracting MNIST_data/train-labels-idx1-ubyte.gz\n",
      "Extracting MNIST_data/t10k-images-idx3-ubyte.gz\n",
      "Extracting MNIST_data/t10k-labels-idx1-ubyte.gz\n"
     ]
    }
   ],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "from itertools import product\n",
    "from math import ceil\n",
    "from utils import *\n",
    "\n",
    "from tensorflow.examples.tutorials.mnist import input_data\n",
    "mnist = input_data.read_data_sets('MNIST_data')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class MNISTEncoder(object):\n",
    "    def __init__(self, emb_dim, sess):\n",
    "        self.emb_dim = emb_dim\n",
    "        self.sess = sess\n",
    "        \n",
    "        self.inputs = tf.placeholder(\"float32\", [None, 28*28])\n",
    "        self.labels = tf.placeholder(\"bool\", [None])\n",
    "        \n",
    "        self._build_model()\n",
    "\n",
    "    def _build_model(self):\n",
    "        # Convolutional encoder\n",
    "        x_image = tf.reshape(self.inputs, [-1,28,28,1])\n",
    "\n",
    "        self.W_conv1 = weight_variable([5, 5, 1, 32])\n",
    "        self.b_conv1 = bias_variable([32])\n",
    "        self.h_conv1 = tf.nn.relu(conv2d(x_image, self.W_conv1) + self.b_conv1)\n",
    "        self.h_pool1 = max_pool_2x2(self.h_conv1)\n",
    "\n",
    "        self.W_conv2 = weight_variable([5, 5, 32, 64])\n",
    "        self.b_conv2 = bias_variable([64])\n",
    "\n",
    "        self.h_conv2 = tf.nn.relu(conv2d(self.h_pool1, self.W_conv2) + self.b_conv2)\n",
    "        self.h_pool2 = max_pool_2x2(self.h_conv2)\n",
    "\n",
    "        self.W_fc1 = weight_variable([7 * 7 * 64, self.emb_dim])\n",
    "        self.b_fc1 = bias_variable([self.emb_dim])\n",
    "\n",
    "        h_pool2_flat = tf.reshape(self.h_pool2, [-1, 7*7*64])\n",
    "        self.emb = tf.matmul(h_pool2_flat, self.W_fc1) + self.b_fc1\n",
    "\n",
    "        # L2 normalize\n",
    "        self.norm_emb = tf.nn.l2_normalize(self.emb, 1)\n",
    "\n",
    "    def get_norm_embedding(self, batch):\n",
    "        return self.sess.run(self.norm_emb, feed_dict={self.inputs: batch})        \n",
    "        \n",
    "    def get_embedding(self, batch):\n",
    "        return self.sess.run(self.emb, feed_dict={self.inputs: batch})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def magnet_loss_old(r, m, d, alpha=1.0):\n",
    "    \"\"\"Compute magnet loss for batch.\n",
    "    \n",
    "    Given a batch of features r consisting of m batches\n",
    "    each with d assigned examples and a cluster separation\n",
    "    gap of alpha, compute the total magnet loss and the per\n",
    "    example losses.\n",
    "    \n",
    "    Args:\n",
    "        r: A batch of features.\n",
    "        m: The number of clusters in the batch.\n",
    "        d: The number of examples in each cluster.\n",
    "        alpha: The cluster separation gap hyperparameter.\n",
    "        \n",
    "    Returns:\n",
    "        total_loss: The total magnet loss for the batch.\n",
    "        losses: The loss for each example in the batch.\n",
    "    \n",
    "    \"\"\"\n",
    "\n",
    "    # Take cluster means within the batch\n",
    "    cluster_means = tf.reduce_mean(tf.reshape(r, [m, d, -1]), 1)\n",
    "\n",
    "    # Compute squared differences of each example to each cluster centroid\n",
    "    sample_cluster_pair_inds = np.array(list(product(range(m*d), range(m))))\n",
    "    sample_costs = tf.squared_difference(\n",
    "        tf.gather(r, sample_cluster_pair_inds[:,0]),\n",
    "        tf.gather(cluster_means, sample_cluster_pair_inds[:,1]))\n",
    "\n",
    "    # Sum to compute squared distances of each example to each cluster centroid\n",
    "    # and reshape such that tensor is indexed by\n",
    "    # [true cluster, comparison cluster, example in true cluster]\n",
    "    sample_costs = tf.reshape(tf.reduce_sum(sample_costs, 1), [m, d, m])\n",
    "    sample_costs = tf.transpose(sample_costs, [0, 2, 1])\n",
    "\n",
    "    # Select distances of examples to their own centroid\n",
    "    same_cluster_inds = np.vstack(np.diag_indices(m)).T\n",
    "    intra_cluster_costs = tf.gather_nd(sample_costs, same_cluster_inds)\n",
    "\n",
    "    # Select distances of examples to other centroids and reshape such that\n",
    "    # tensor is indexed by [true cluster, comparison cluster, example]\n",
    "    cluster_inds = np.arange(m)\n",
    "    diff_cluster_inds = np.vstack(\n",
    "        [np.repeat(cluster_inds, m-1), \n",
    "         np.hstack([cluster_inds[cluster_inds != i] for i in range(m)])]).T\n",
    "    inter_cluster_costs = tf.reshape(tf.gather_nd(sample_costs, diff_cluster_inds), [m, m-1, d])\n",
    "\n",
    "    # Compute variance of intra-cluster squared distances\n",
    "    variance = tf.reduce_sum(intra_cluster_costs) / (m * d - 1)\n",
    "    var_normalizer = -1 / 2*variance**2\n",
    "\n",
    "    # Compute numerator and denominator of inner term\n",
    "    numerator = tf.exp(var_normalizer * intra_cluster_costs - alpha)\n",
    "    denominator = tf.reduce_sum(tf.exp(var_normalizer * inter_cluster_costs), 1)\n",
    "\n",
    "    # Compute example losses and total loss\n",
    "    losses = tf.nn.relu(-tf.log(numerator / denominator))\n",
    "    total_loss = tf.reduce_mean(losses)\n",
    "    \n",
    "    return total_loss, losses\n",
    "\n",
    "\n",
    "def magnet_loss(r, m, d, alpha=1.0):\n",
    "    \"\"\"Compute magnet loss for batch.\n",
    "    \n",
    "    Given a batch of features r consisting of m batches\n",
    "    each with d assigned examples and a cluster separation\n",
    "    gap of alpha, compute the total magnet loss and the per\n",
    "    example losses.\n",
    "    \n",
    "    Args:\n",
    "        r: A batch of features.\n",
    "        m: The number of clusters in the batch.\n",
    "        d: The number of examples in each cluster.\n",
    "        alpha: The cluster separation gap hyperparameter.\n",
    "        \n",
    "    Returns:\n",
    "        total_loss: The total magnet loss for the batch.\n",
    "        losses: The loss for each example in the batch.\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    # Helper to compute indexes to select intra- and inter-cluster\n",
    "    # distances\n",
    "    def compute_comparison_inds():\n",
    "        same_cluster_inds = []\n",
    "        for i in range(m*d):\n",
    "            c = i / d\n",
    "            same_cluster_inds.append(c*d*m + c*d + (i % d))\n",
    "        diff_cluster_inds = sorted(set(range(m*m*d)) - set(same_cluster_inds))\n",
    "        \n",
    "        return same_cluster_inds, diff_cluster_inds\n",
    "\n",
    "        \n",
    "    # Take cluster means within the batch\n",
    "    cluster_means = tf.reduce_mean(tf.reshape(r, [m, d, -1]), 1)\n",
    "\n",
    "    # Compute squared distance of each example to each cluster centroid\n",
    "    sample_cluster_pair_inds = np.array(list(product(range(m), range(m*d))))\n",
    "    sample_costs = tf.squared_difference(\n",
    "        tf.gather(cluster_means, sample_cluster_pair_inds[:,0]),\n",
    "        tf.gather(r, sample_cluster_pair_inds[:,1]))\n",
    "    sample_costs = tf.reduce_sum(sample_costs, 1)\n",
    "    \n",
    "    # Compute intra- and inter-cluster comparison indexes\n",
    "    same_cluster_inds, diff_cluster_inds = compute_comparison_inds()\n",
    "    \n",
    "    # Select distances of examples to their own centroid\n",
    "    intra_cluster_costs = tf.gather(sample_costs, same_cluster_inds)\n",
    "    intra_cluster_costs = tf.reshape(intra_cluster_costs, [m, d])\n",
    "    \n",
    "    # Select distances of examples to other centroids\n",
    "    inter_cluster_costs = tf.reshape(tf.gather(sample_costs, diff_cluster_inds), [m, m-1, d])\n",
    "\n",
    "    # Compute variance of intra-cluster squared distances\n",
    "    variance = tf.reduce_sum(intra_cluster_costs) / (m * d - 1)\n",
    "    var_normalizer = -1 / 2*variance**2\n",
    "\n",
    "    # Compute numerator and denominator of inner term\n",
    "    numerator = tf.exp(var_normalizer * intra_cluster_costs - alpha)\n",
    "    denominator = tf.reduce_sum(tf.exp(var_normalizer * inter_cluster_costs), 1)\n",
    "\n",
    "    # Compute example losses and total loss\n",
    "    losses = tf.nn.relu(-tf.log(numerator / denominator))\n",
    "    total_loss = tf.reduce_mean(losses)\n",
    "    \n",
    "    return total_loss, losses\n",
    "\n",
    "\n",
    "    \n",
    "# sess = tf.InteractiveSession()\n",
    "\n",
    "# m = 3\n",
    "# d = 2\n",
    "\n",
    "# m = 6\n",
    "# d = 4\n",
    "\n",
    "# K = 5\n",
    "# alpha = 15.0\n",
    "\n",
    "# r = tf.placeholder(tf.float32, [None, 8])\n",
    "# magnet_loss1, losses1 = magnet_loss_old(r, m, d, alpha)\n",
    "# magnet_loss2, losses2 = magnet_loss(r, m, d, alpha)\n",
    "\n",
    "\n",
    "# # Helper to generate debug data\n",
    "# def gen_data(m, d):\n",
    "#     data = []\n",
    "#     for c in range(m):\n",
    "#         a = (c + 1) * 3\n",
    "#         centroid = np.random.random([1, 8]) * a\n",
    "#         data.append(centroid + np.random.random([d, 8]))\n",
    "#     return np.vstack(data)\n",
    "\n",
    "\n",
    "# # feed_dict = {r: np.random.random([m*d, 8])}\n",
    "# feed_dict = {r: gen_data(m, d)}\n",
    "\n",
    "# print sess.run([magnet_loss1, magnet_loss2], feed_dict=feed_dict)\n",
    "\n",
    "\n",
    "# sess.close()\n",
    "# tf.reset_default_graph()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensor(\"StridedSlice:0\", shape=(6,), dtype=int32)\n",
      "[1.8370082, array([0, 1, 2, 3, 4, 5], dtype=int32)]\n"
     ]
    }
   ],
   "source": [
    "\n",
    "def magnet_loss(r, c, m, d, alpha=1.0):\n",
    "    \"\"\"Compute magnet loss for batch.\n",
    "    \n",
    "    Given a batch of features r consisting of m batches\n",
    "    each with d assigned examples and a cluster separation\n",
    "    gap of alpha, compute the total magnet loss and the per\n",
    "    example losses.\n",
    "    \n",
    "    Args:\n",
    "        r: A batch of features.\n",
    "        c: Class labels for each example.\n",
    "        m: The number of clusters in the batch.\n",
    "        d: The number of examples in each cluster.\n",
    "        alpha: The cluster separation gap hyperparameter.\n",
    "        \n",
    "    Returns:\n",
    "        total_loss: The total magnet loss for the batch.\n",
    "        losses: The loss for each example in the batch.\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    # Helper to compute indexes to select intra- and inter-cluster\n",
    "    # distances\n",
    "    def compute_comparison_inds():\n",
    "        same_cluster_inds = []\n",
    "        for i in range(m*d):\n",
    "            c = i / d\n",
    "            same_cluster_inds.append(c*d*m + c*d + (i % d))\n",
    "        diff_cluster_inds = sorted(set(range(m*m*d)) - set(same_cluster_inds))\n",
    "        \n",
    "        return same_cluster_inds, diff_cluster_inds\n",
    "\n",
    "        \n",
    "    # Take cluster means within the batch\n",
    "    cluster_means = tf.reduce_mean(tf.reshape(r, [m, d, -1]), 1)\n",
    "\n",
    "    # Compute squared distance of each example to each cluster centroid\n",
    "    sample_cluster_pair_inds = np.array(list(product(range(m), range(m*d))))\n",
    "    sample_costs = tf.squared_difference(\n",
    "        tf.gather(cluster_means, sample_cluster_pair_inds[:,0]),\n",
    "        tf.gather(r, sample_cluster_pair_inds[:,1]))\n",
    "    sample_costs = tf.reduce_sum(sample_costs, 1)\n",
    "    \n",
    "    # Compute intra- and inter-cluster comparison indexes\n",
    "    same_cluster_inds, diff_cluster_inds = compute_comparison_inds()\n",
    "    \n",
    "    # Select distances of examples to their own centroid\n",
    "    intra_cluster_costs = tf.gather(sample_costs, same_cluster_inds)\n",
    "    intra_cluster_costs = tf.reshape(intra_cluster_costs, [m, d])\n",
    "    \n",
    "    \n",
    "    \n",
    "    # Select distances of examples to other centroids\n",
    "    inter_cluster_costs = tf.reshape(tf.gather(sample_costs, diff_cluster_inds), [m, m-1, d])\n",
    "    \n",
    "    \n",
    "    # Select distances of examples to other class centroids\n",
    "    cluster_classes = tf.strided_slice(c, [0], [m*d], [d])\n",
    "    print cluster_classes\n",
    "    \n",
    "\n",
    "    # Compute variance of intra-cluster squared distances\n",
    "    variance = tf.reduce_sum(intra_cluster_costs) / (m * d - 1)\n",
    "    var_normalizer = -1 / 2*variance**2\n",
    "\n",
    "    # Compute numerator and denominator of inner term\n",
    "    numerator = tf.exp(var_normalizer * intra_cluster_costs - alpha)\n",
    "    denominator = tf.reduce_sum(tf.exp(var_normalizer * inter_cluster_costs), 1)\n",
    "\n",
    "    # Compute example losses and total loss\n",
    "    losses = tf.nn.relu(-tf.log(numerator / denominator))\n",
    "    total_loss = tf.reduce_mean(losses)\n",
    "    \n",
    "    return total_loss, losses\n",
    "\n",
    "\n",
    "\n",
    "sess = tf.InteractiveSession()\n",
    "\n",
    "m = 6\n",
    "d = 4\n",
    "\n",
    "K = 5\n",
    "alpha = 15.0\n",
    "\n",
    "r = tf.placeholder(tf.float32, [m*d, 8])\n",
    "c = tf.placeholder(tf.int32, [m*d])\n",
    "magnet_loss2, losses2, cc = magnet_loss(r, c, m, d, alpha)\n",
    "\n",
    "\n",
    "# Helper to generate debug data\n",
    "def gen_data(m, d):\n",
    "    data = []\n",
    "    for c in range(m):\n",
    "        a = (c + 1) * 3\n",
    "        centroid = np.random.random([1, 8]) * a\n",
    "        data.append(centroid + np.random.random([d, 8]))\n",
    "    return np.vstack(data)\n",
    "\n",
    "\n",
    "# feed_dict = {r: np.random.random([m*d, 8])}\n",
    "feed_dict = {r: gen_data(m, d), c: np.repeat(range(m), d)}\n",
    "\n",
    "print sess.run([magnet_loss2, cc], feed_dict=feed_dict)\n",
    "\n",
    "\n",
    "sess.close()\n",
    "tf.reset_default_graph()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5,\n",
       "       5])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.repeat(range(m), d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from sklearn.cluster import KMeans\n",
    "\n",
    "class ClusterBatchBuilder(object):\n",
    "    def __init__(self, labels, k, m, d):\n",
    "        \n",
    "        self.num_classes = np.unique(labels).shape[0]\n",
    "        self.labels = labels\n",
    "\n",
    "        self.k = k\n",
    "        self.m = m\n",
    "        self.d = d\n",
    "\n",
    "        self.centroids = None\n",
    "        self.assignments = np.zeros_like(labels, int)\n",
    "        self.cluster_assignments = {}\n",
    "        self.example_losses = None\n",
    "        self.cluster_losses = np.zeros([self.k * self.num_classes], float)\n",
    "\n",
    "    \n",
    "    def update_clusters(self, rep_data, max_iter=20):\n",
    "        \"\"\"\n",
    "        Given an array of representations for the entire training set,\n",
    "        recompute clusters and store example cluster assignments in a\n",
    "        quickly sampleable form.\n",
    "        \"\"\"\n",
    "        # Lazily allocate array for centroids\n",
    "        if self.centroids is None:\n",
    "            self.centroids = np.zeros([self.num_classes * self.k, rep_data.shape[1]])\n",
    "        \n",
    "        for c in range(self.num_classes):\n",
    "\n",
    "            class_mask = self.labels == c\n",
    "            class_examples = rep_data[class_mask]\n",
    "            kmeans = KMeans(n_clusters=self.k, init='k-means++', n_init=1, max_iter=max_iter)\n",
    "            kmeans.fit(class_examples)\n",
    "\n",
    "            # Save cluster centroids for finding impostor clusters\n",
    "            start = self.get_cluster_ind(c, 0)\n",
    "            stop = self.get_cluster_ind(c, self.k)\n",
    "            self.centroids[start:stop] = kmeans.cluster_centers_\n",
    "\n",
    "            # Update assignments with new global cluster indexes\n",
    "            self.assignments[class_mask] = self.get_cluster_ind(c, kmeans.predict(class_examples))\n",
    "            \n",
    "        # Construct a map from cluster to example indexes for fast batch creation\n",
    "        for cluster in range(self.k * self.num_classes):\n",
    "            cluster_mask = self.assignments == cluster\n",
    "            self.cluster_assignments[cluster] = np.flatnonzero(cluster_mask)\n",
    "\n",
    "        \n",
    "    def update_losses(self, indexes, losses):\n",
    "        \"\"\"\n",
    "        Given a list of examples indexes and corresponding losses\n",
    "        store the new losses and update corresponding cluster losses.\n",
    "        \"\"\"\n",
    "        \n",
    "        # Update example losses\n",
    "        self.example_losses[indexes] = losses\n",
    "\n",
    "        # Find affected clusters and update the corresponding cluster losses\n",
    "        clusters = np.unique(self.assignments[indexes])\n",
    "        for cluster in clusters:\n",
    "            cluster_example_losses = self.example_losses[self.assignments == cluster]\n",
    "            self.cluster_losses[cluster] = np.mean(cluster_example_losses)\n",
    "\n",
    "        \n",
    "    def gen_batch(self):\n",
    "        \"\"\"\n",
    "        Sample a batch by first sampling a seed cluster proportionally to\n",
    "        the mean loss of the clusters, then finding nearest neighbor\n",
    "        \"impostor\" clusters, then sampling d examples uniformly from each cluster.\n",
    "        \n",
    "        The generated batch will consist of m clusters each with d consecutive\n",
    "        examples.\n",
    "        \"\"\"\n",
    "        \n",
    "        # Sample seed cluster proportionally to cluster losses if available\n",
    "        if self.cluster_losses is not None:\n",
    "            p = self.cluster_losses / np.sum(self.cluster_losses)\n",
    "            seed_cluster = np.random.choice(self.num_classes * self.k, p=p)\n",
    "        else:\n",
    "            seed_cluster = np.random.choice(self.num_classes * self.k)\n",
    "\n",
    "        # Get imposter clusters by ranking centroids by distance\n",
    "        # The seed cluster itself is guaranteed to be included\n",
    "        sq_dists = ((self.centroids[seed_cluster] - self.centroids) ** 2).sum(axis=1)\n",
    "        clusters = np.argpartition(sq_dists, self.m)[:self.m]\n",
    "        \n",
    "        # Sample examples uniformly from cluster\n",
    "        batch_indexes = np.empty([self.m * self.d], int)\n",
    "        for i, c in enumerate(clusters):\n",
    "            x = np.random.choice(self.cluster_assignments[c], self.d, replace=False)\n",
    "\n",
    "            start = i * self.d\n",
    "            stop = start + self.d\n",
    "            batch_indexes[start:stop] = x\n",
    "\n",
    "        return batch_indexes\n",
    "\n",
    "    \n",
    "    def get_cluster_ind(self, c, i):\n",
    "        \"\"\"\n",
    "        Given a class index and a cluster index within the class\n",
    "        return the global cluster index\n",
    "        \"\"\"\n",
    "        return c * self.k + i\n",
    "    \n",
    "\n",
    "# rep = np.random.random([mnist.train.images.shape[0], 8])\n",
    "# batch_builder = ClusterBatchBuilder(mnist.train.labels, 5, 12, 4)\n",
    "# batch_builder.update_clusters(rep)\n",
    "# batch_builder.update_losses(range(mnist.train.labels.shape[0]), np.random.random([mnist.train.labels.shape[0]]))\n",
    "# batch_builder.gen_batch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Define magnet loss parameters\n",
    "m = 6\n",
    "d = 4\n",
    "k = 3\n",
    "alpha = 1.0\n",
    "batch_size = m * d\n",
    "\n",
    "# Define training data\n",
    "X = mnist.train.images\n",
    "y = mnist.train.labels\n",
    "\n",
    "# Define model and training parameters\n",
    "epoch_steps = int(ceil(float(X.shape[0]) / batch_size)) \n",
    "emb_dim = 2\n",
    "# num_steps = 5 * epoch_steps\n",
    "# cluster_refresh_interval = epoch_steps\n",
    "num_steps = epoch_steps\n",
    "cluster_refresh_interval = 20\n",
    "\n",
    "\n",
    "\n",
    "sess = tf.InteractiveSession()\n",
    "\n",
    "# Model\n",
    "with tf.variable_scope('model'):\n",
    "    model = MNISTEncoder(emb_dim, sess)\n",
    "\n",
    "# Loss\n",
    "with tf.variable_scope('magnet_loss'):\n",
    "    train_loss, losses = magnet_loss(model.emb, m, d, alpha)\n",
    "\n",
    "train_op = tf.train.AdamOptimizer(1e-4).minimize(train_loss)\n",
    "\n",
    "sess.run(tf.initialize_all_variables())\n",
    "\n",
    "\n",
    "\n",
    "def compute_reps(extract_fn, X, chunk_size):\n",
    "    chunks = int(ceil(float(X.shape[0]) / chunk_size))\n",
    "    reps = []\n",
    "    for i in range(chunks):\n",
    "        start = i * chunk_size\n",
    "        stop = start + chunk_size\n",
    "        chunk_reps = extract_fn(X[start:stop])\n",
    "        reps.append(chunk_reps)\n",
    "    return np.vstack(reps)\n",
    "\n",
    "# Get initial embedding\n",
    "extract = lambda x: sess.run(model.emb, feed_dict={model.inputs: x})\n",
    "initial_reps = compute_reps(extract, X, 400)\n",
    "\n",
    "print initial_reps.shape\n",
    "\n",
    "\n",
    "# # Create batcher\n",
    "batch_builder = ClusterBatchBuilder(mnist.train.labels, k, m, d)\n",
    "batch_builder.update_clusters(initial_reps)\n",
    "\n",
    "\n",
    "for i in range(num_steps):\n",
    "    \n",
    "    batch_inds = batch_builder.gen_batch()\n",
    "    _, batch_loss = sess.run([train_op, train_loss], feed_dict={model.inputs: X[batch_inds]})\n",
    "    \n",
    "    if not i % 100:\n",
    "        print i, batch_loss\n",
    "    \n",
    "    if not i % cluster_refresh_interval:\n",
    "        print 'Refreshing clusters'\n",
    "        reps = compute_reps(extract, X, 400)\n",
    "        print reps[0]\n",
    "        print np.any(np.isnan(reps))\n",
    "        batch_builder.update_clusters(reps)\n",
    "        \n",
    "final_reps = compute_reps(extract, X, 400)\n",
    "    \n",
    "sess.close()\n",
    "tf.reset_default_graph()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "reps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "num_plot = 500\n",
    "imgs = mnist.train.images[:num_plot]\n",
    "imgs = np.reshape(imgs, [num_plot, 28, 28])\n",
    "plot_embedding(initial_reps[:num_plot], mnist.train.labels[:num_plot], imgs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "num_plot = 500\n",
    "imgs = mnist.train.images[:num_plot]\n",
    "imgs = np.reshape(imgs, [num_plot, 28, 28])\n",
    "plot_embedding(final_reps[:num_plot], mnist.train.labels[:num_plot], imgs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "sess.close()\n",
    "tf.reset_default_graph()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}