Question

亲爱的TensorFlow社区，

我正在使用tf.contrib.factorization.KMeansClustering训练分类器，但是培训的进度确实很慢，只使用了我GPU的1％。

但是，我的4个CPU内核的使用率一直在35％左右。

是不是为CPU写的K-Means数比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)

最好，乔什

Answer 1

这是迄今为止我对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

Answer 2

我发现增加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

的最佳实践指南

TensorFlow中更快的K-Means聚类

2 个答案: