title: Faster K-Means Clustering in TensorFlow
body: <p>Dear TensorFlow Community,</p>
  
  <p>I'm training a classifier with <a href="https://www.tensorflow.org/api_docs/python/tf/contrib/factorization/KMeansClustering" rel="noreferrer"><code>tf.contrib.factorization.KMeansClustering</code></a>,
  but the training goes really slow, and only uses 1% of my GPU.</p>
  
  <p>However, my 4 CPU cores are hitting about 35% use constantly.</p>
  
  <p>Is it the case that K-Means is written more for the CPU than the GPU?</p>
  
  <p>Is there a way I can shift more of the computation to the GPU, or some
  other approach to speed up training?</p>
  
  <p>Below is my script for training (Python3).</p>
  
  <p>Thank you for your time.</p>
  
  <pre><code>import tensorflow as tf 
  
  
  
  def parser(record):
  
    features={
      'feats': tf.FixedLenFeature([], tf.string),
    }
  
    parsed = tf.parse_single_example(record, features)
    feats = tf.convert_to_tensor(tf.decode_raw(parsed['feats'], tf.float64))
  
    return {'feats': feats}
  
  
  def my_input_fn(tfrecords_path):
  
      dataset = (
          tf.data.TFRecordDataset(tfrecords_path)
          .map(parser)
          .batch(1024)
      )
  
      iterator = dataset.make_one_shot_iterator()
      batch_feats = iterator.get_next()
  
      return batch_feats
  
  
  ### SPEC FUNCTIONS ###
  
  train_spec_kmeans = tf.estimator.TrainSpec(input_fn = lambda: my_input_fn('/home/ubuntu/train.tfrecords') , max_steps=10000)
  eval_spec_kmeans = tf.estimator.EvalSpec(input_fn = lambda: my_input_fn('/home/ubuntu/eval.tfrecords') )
  
  
  
  ### INIT ESTIMATOR ###
  
  KMeansEstimator = tf.contrib.factorization.KMeansClustering(
      num_clusters=500,
      feature_columns = [tf.feature_column.numeric_column(
          key='feats',
          dtype=tf.float64,
          shape=(377,),
      )],
      use_mini_batch=True)
  
  
  ### TRAIN &amp; EVAL ###
  
  tf.estimator.train_and_evaluate(KMeansEstimator, train_spec_kmeans, eval_spec_kmeans)
  </code></pre>
  
  <p>Best,
  Josh</p>