TF KMeansClustering 在 GPU 上没有 运行

TF KMeansClustering doesn't run on GPU

运行 Ubuntu 16.04,最新 (1.1.0) tensorflow(通过 pip3 install tensorflow-gpu 安装),CUDA8 + CUDNN5。

代码大致如下所示:

import tensorflow as tf
from tensorflow.contrib.learn import KMeansClustering

trainencflt = #pandas frame with ~30k rows and ~300 columns
def train_input_fn():
    return (tf.constant(trainencflt, shape = [trainencflt.shape[0], trainencflt.shape[1]]), None)

configuration = tf.contrib.learn.RunConfig(log_device_placement=True)
model = KMeansClustering(num_clusters=k,
                         initial_clusters=KMeansClustering.RANDOM_INIT,
                         relative_tolerance=1e-8,
                         config=configuration)
model.fit(input_fn = train_input_fn, steps = 100)

运行时我看到:

2017-06-15 10:24:41.564890: I tensorflow/core/common_runtime/gpu/gpu_device.cc:887] Found device 0 with properties:
name: GeForce GTX 1080
major: 6 minor: 1 memoryClockRate (GHz) 1.7335
pciBusID 0000:81:00.0
Total memory: 7.92GiB
Free memory: 7.81GiB
2017-06-15 10:24:41.564934: I tensorflow/core/common_runtime/gpu/gpu_device.cc:908] DMA: 0
2017-06-15 10:24:41.564942: I tensorflow/core/common_runtime/gpu/gpu_device.cc:918] 0:   Y
2017-06-15 10:24:41.564956: I tensorflow/core/common_runtime/gpu/gpu_device.cc:977] Creating TensorFlow device (/gpu:0) -> (device: 0, name: GeForce GTX 1080, pci bus id: 0000:81:00.0)

内存分配:

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID  Type  Process name                               Usage      |
|=============================================================================|
|    1       548    C   python                                        7745MiB |
+-----------------------------------------------------------------------------+

但是 none 的操作是在 GPU 上执行的(它一直保持在 0%,CPU 利用率在所有内核上都飙升):

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 375.66                 Driver Version: 375.66                    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   1  GeForce GTX 1080    Off  | 0000:02:00.0     Off |                  N/A |
| 29%   43C    P8    13W / 180W |   7747MiB /  8114MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

没有看到任何放置日志(即使我将 log_device_placement 指定为 True)。

我确实尝试了 the simple GPU examples,它们工作得很好(至少放置日志看起来不错)。

我是不是漏掉了什么?

查看代码库 - TF 1.1.0 根本没有 GPU 内核。