使用 tensorflow KMeans 进行聚类

数据

为了方便, 我们使用 MNIST 数据集.

核心 API

(all_scores, cluster_idx, scores, cluster_centers_initialized, init_op, train_op) = kmeans.training_graph()

这个 API 用来生成一个 kmeans 算法的计算图.

API 返回一个元组:

  • all_scores: 返回一个维度为 (num_input, num_clusters) 的矩阵, 值是每个 input vector 和 cluster center 的距离.
  • cluster_idx: 一个 vector, 包含了 cluster 的 id. TODO inp ?
  • scores: 对应与每个 cluster 的距离.
  • cluster_centers_initialized: 标记是否被初始化.
  • init_op: 初始化 clusters 的 op.
  • train_op: 训练用的的 op.

开始

首先加载 MNIST 数据:


# 加载 MNIST 数据
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("/tmp/data/", one_hot=True)
full_data_x = mnist.train.images

# 初始化参数
num_steps = 50
batch_size = 1024
k = 25
num_classes = 10
num_features = 784 # 28x28 像素

然后是构建训练的计算图 :

# 输入
X = tf.placeholder(tf.float32, shape=[None, num_features])

# KMeans 计算图
kmeans = KMeans(inputs=X, num_clusters=k, distance_metric='cosine',
use_mini_batch=True)
(all_scores, cluster_idx, scores, cluster_centers_initialized,
init_op, train_op) = kmeans.training_graph()
cluster_idx = cluster_idx[0]
avg_distance = tf.reduce_mean(scores)

# 初始化
init_vars = tf.global_variables_initializer()
sess = tf.Session()
sess.run(init_vars, feed_dict={X: full_data_x})
sess.run(init_op, feed_dict={X: full_data_x})

# 训练
for i in range(1, num_steps + 1):
# idx, shape: (55000,), 值为 0-k
_, d, idx = sess.run([train_op, avg_distance, cluster_idx],
feed_dict={X: full_data_x})
if i % 10 == 0 or i == 1:
print("Step %i, Avg Distance: %f" % (i, d))

运行后, 可以查看训练的输出:

Step 1, Avg Distance: 0.341471
Step 10, Avg Distance: 0.221609
Step 20, Avg Distance: 0.220328
Step 30, Avg Distance: 0.219776
Step 40, Avg Distance: 0.219419
Step 50, Avg Distance: 0.219154

可视化

为了直观的看结果的分布, 我们可以用散点图来显示, 先使用 PCA 再根据对应 label 画散点图.

# counts, shape: (25, 10)
counts = np.zeros(shape=(k, num_classes))
for i in range(len(idx)):
counts[idx[i]] += mnist.train.labels[i]
# 各个 cluster 的 class
labels_map = [np.argmax(c) for c in counts] # (25,10)
labels_map = tf.convert_to_tensor(labels_map)
cluster_label = tf.nn.embedding_lookup(labels_map, cluster_idx)

plot:

from sklearn.decomposition import PCA
pca = PCA(n_components=2).fit(mnist.train.images)
pca_2d = pca.transform(mnist.train.images)
import matplotlib.pyplot as plt
plt.figure(1)
plt.clf()

# 分类颜色
colors = np.random.random((num_classes, 3))
classes = sess.run(cluster_label, feed_dict={X: full_data_x})
colors_of_classes = [colors[c] for c in classes]
plt.scatter(pca_2d[:, 0], pca_2d[:, 1], c=colors_of_classes, s=1)
plt.show()

将训练后的结果 plot 之后, 还是分类比较明显的 (虽然准确度不高...):

其他

1.4 的 API 已经有返回 cluster 中心了, 请看这里

cluster_centers_var: a Variable holding the cluster centers.

参考


https://github.com/aymericdamien/TensorFlow-Examples/blob/master/examples/2_BasicModels/kmeans.py
http://scikit-learn.org/stable/auto_examples/cluster/plot_kmeans_digits.html