kmeans.py 2.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687
  1. """ K-Means.
  2. Implement K-Means algorithm with TensorFlow, and apply it to classify
  3. handwritten digit images. This example is using the MNIST database of
  4. handwritten digits as training samples (http://yann.lecun.com/exdb/mnist/).
  5. Note: This example requires TensorFlow v1.1.0 or over.
  6. Author: Aymeric Damien
  7. Project: https://github.com/aymericdamien/TensorFlow-Examples/
  8. """
  9. from __future__ import print_function
  10. import numpy as np
  11. import tensorflow as tf
  12. from tensorflow.contrib.factorization import KMeans
  13. # Ignore all GPUs, tf random forest does not benefit from it.
  14. import os
  15. os.environ["CUDA_VISIBLE_DEVICES"] = ""
  16. # Import MNIST data
  17. from tensorflow.examples.tutorials.mnist import input_data
  18. mnist = input_data.read_data_sets("/tmp/data/", one_hot=True)
  19. full_data_x = mnist.train.images
  20. # Parameters
  21. num_steps = 50 # Total steps to train
  22. batch_size = 1024 # The number of samples per batch
  23. k = 25 # The number of clusters
  24. num_classes = 10 # The 10 digits
  25. num_features = 784 # Each image is 28x28 pixels
  26. # Input images
  27. X = tf.placeholder(tf.float32, shape=[None, num_features])
  28. # Labels (for assigning a label to a centroid and testing)
  29. Y = tf.placeholder(tf.float32, shape=[None, num_classes])
  30. # K-Means Parameters
  31. kmeans = KMeans(inputs=X, num_clusters=k, distance_metric='cosine',
  32. use_mini_batch=True)
  33. # Build KMeans graph
  34. (all_scores, cluster_idx, scores, cluster_centers_initialized, init_op,
  35. train_op) = kmeans.training_graph()
  36. cluster_idx = cluster_idx[0] # fix for cluster_idx being a tuple
  37. avg_distance = tf.reduce_mean(scores)
  38. # Initialize the variables (i.e. assign their default value)
  39. init_vars = tf.global_variables_initializer()
  40. # Start TensorFlow session
  41. sess = tf.Session()
  42. # Run the initializer
  43. sess.run(init_vars, feed_dict={X: full_data_x})
  44. sess.run(init_op, feed_dict={X: full_data_x})
  45. # Training
  46. for i in range(1, num_steps + 1):
  47. _, d, idx = sess.run([train_op, avg_distance, cluster_idx],
  48. feed_dict={X: full_data_x})
  49. if i % 10 == 0 or i == 1:
  50. print("Step %i, Avg Distance: %f" % (i, d))
  51. # Assign a label to each centroid
  52. # Count total number of labels per centroid, using the label of each training
  53. # sample to their closest centroid (given by 'idx')
  54. counts = np.zeros(shape=(k, num_classes))
  55. for i in range(len(idx)):
  56. counts[idx[i]] += mnist.train.labels[i]
  57. # Assign the most frequent label to the centroid
  58. labels_map = [np.argmax(c) for c in counts]
  59. labels_map = tf.convert_to_tensor(labels_map)
  60. # Evaluation ops
  61. # Lookup: centroid_id -> label
  62. cluster_label = tf.nn.embedding_lookup(labels_map, cluster_idx)
  63. # Compute accuracy
  64. correct_prediction = tf.equal(cluster_label, tf.cast(tf.argmax(Y, 1), tf.int32))
  65. accuracy_op = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
  66. # Test Model
  67. test_x, test_y = mnist.test.images, mnist.test.labels
  68. print("Test Accuracy:", sess.run(accuracy_op, feed_dict={X: test_x, Y: test_y}))