小编JR *_*yer的帖子

亲爱的TensorFlow社区,

我正在训练分类器tf.contrib.factorization.KMeansClustering,但是训练非常慢,并且只使用了1%的GPU.

但是,我的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, …