尊敬的 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)
最好的,
乔什
最佳答案
到目前为止,这是我的最佳答案 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
关于python-3.x - TensorFlow 中更快的 K-Means 聚类,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/50927298/