TensorFlow 中更快的 K-Means 聚类

Posted 2023-03-12

【中文标题】TensorFlow 中更快的 K-Means 聚类

【英文标题】：Faster K-Means Clustering in TensorFlow

【发布时间】：2018-11-28 08:27:08

【问题描述】：

尊敬的 TensorFlow 社区，

我正在使用tf.contrib.factorization.KMeansClustering 训练分类器， 但训练进行得很慢，而且只使用了我的 GPU 的 1%。

但是，我的 4 个 CPU 内核的使用率持续达到约 35%。

K-Means 是为 CPU 编写的而不是为 GPU 编写的吗？

有没有办法可以将更多的计算转移到 GPU 上，或者一些 其他加快训练速度的方法？

下面是我的训练脚本（Python3）。

感谢您的宝贵时间。

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 & EVAL ###

tf.estimator.train_and_evaluate(KMeansEstimator, train_spec_kmeans, eval_spec_kmeans)

最好， 乔什

【问题讨论】：

将所有tf.float64 更改为tf.float32 看看是否有帮助。

如果我转到tf.float32，我会看到很多NaNs：(

【参考方案1】：

我看到的增加 gpu 和 cpu 使用率的一件事是在数据集上使用预取。它让数据集生产者在模型也在消耗前一批数据的同时获取数据，从而最大限度地利用资源。还指定你的 CPU 的最大值会加快这个过程。 我会这样重组它

dataset = (
    tf.data.TFRecordDataset(tfrecords_path)
    .map(parser,num_parallel_calls=multiprocessing.cpu_count())
    .batch(1024)

)

dataset = dataset.prefetch(1024)

在使用 TfRecords here

时，这是一个很好的最佳实践指南

【讨论】：

您认为将不同的变量（簇数、批次数等）作为 2 的幂有多大帮助？我在1024 有批次，但我没有控制我的集群数量。谢谢！

也来自您链接到的文档：tf.contrib.data.map_and_batch 将并行创建批处理，融合 dataset.map() 和 dataset.batch()

我不太确定它是否会产生影响...但是由于内存单元是 32/64 位，因此拥有多个单元可能是访问内存的更好方法。 .我只是不了解 tf.contrib.data.map_and_batch，从你的测量来看似乎工作得更快。感谢您的提示

【参考方案2】：

这是迄今为止我最好的答案，其中包含 time 信息，以 Eliethesaiyan 的答案和 link to docs 为基础。

我原来的Dataset 代码块和性能：

dataset = (
 tf.data.TFRecordDataset(tfrecords_path)
 .map(parse_fn)
 .batch(1024)
)

real    1m36.171s
user    2m57.756s
sys     0m42.304s

Eliethesaiyan 的回答 (prefetch + num_parallel_calls)

dataset = (
    tf.data.TFRecordDataset(tfrecords_path)
    .map(parse_fn,num_parallel_calls=multiprocessing.cpu_count())
    .batch(1024)
    .prefetch(1024)
   )

real  0m41.450s
user  1m33.120s
sys   0m18.772s

来自使用map_and_batch + num_parallel_batches + prefetch 的文档：

dataset = (
    tf.data.TFRecordDataset(tfrecords_path)
    .apply(
       tf.contrib.data.map_and_batch(
          map_func=parse_fn,
          batch_size=1024,
          num_parallel_batches=multiprocessing.cpu_count()
        )
    )
    .prefetch(1024)
 )

real   0m32.855s
user   1m11.412s
sys    0m10.408s