Python 即使指定了其他gpu，skflow也会在gpu0中分配内存

我在4 GPU Amazon实例上遇到了这个问题，使用了一个简单的示例脚本：

import skflow
import tensorflow as tf
from sklearn import datasets

iris = datasets.load_iris()
X_train, X_test, y_train, y_test = cross_validation.train_test_split(iris.data, iris.target,
    test_size=0.2, random_state=42)

def my_model(X, y):

    with tf.device('/gpu:1'):
        layers = skflow.ops.dnn(X, [1000, 500, 150], keep_prob=0.5) # many neurons to see the impac on memory
    with tf.device('/cpu:0'):
        return skflow.models.logistic_regression(layers, y)

classifier = skflow.TensorFlowEstimator(model_fn=my_model, n_classes=3)
classifier.fit(X_train, y_train)

在启动脚本之前，

nvidia smi

的结果是：

Fri Feb 19 11:30:22 2016       
+------------------------------------------------------+                       
| NVIDIA-SMI 346.46     Driver Version: 346.46         |                       
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GRID K520           Off  | 0000:00:03.0     Off |                  N/A |
| N/A   40C    P0    41W / 125W |   2247MiB /  4095MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   1  GRID K520           Off  | 0000:00:04.0     Off |                  N/A |
| N/A   36C    P0    40W / 125W |   2113MiB /  4095MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   2  GRID K520           Off  | 0000:00:05.0     Off |                  N/A |
| N/A   41C    P0    43W / 125W |     53MiB /  4095MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   3  GRID K520           Off  | 0000:00:06.0     Off |                  N/A |
| N/A   39C    P0    41W / 125W |   1816MiB /  4095MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

在脚本运行时：

Fri Feb 19 11:30:53 2016       
+------------------------------------------------------+                       
| NVIDIA-SMI 346.46     Driver Version: 346.46         |                       
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GRID K520           Off  | 0000:00:03.0     Off |                  N/A |
| N/A   40C    P0    46W / 125W |   3926MiB /  4095MiB |     26%      Default |
+-------------------------------+----------------------+----------------------+
|   1  GRID K520           Off  | 0000:00:04.0     Off |                  N/A |
| N/A   37C    P0    42W / 125W |   3926MiB /  4095MiB |     17%      Default |
+-------------------------------+----------------------+----------------------+
|   2  GRID K520           Off  | 0000:00:05.0     Off |                  N/A |
| N/A   41C    P0    44W / 125W |     92MiB /  4095MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   3  GRID K520           Off  | 0000:00:06.0     Off |                  N/A |
| N/A   39C    P0    42W / 125W |   1856MiB /  4095MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

---

因此，内存被分配给GPU0，即使代码中没有任何部分提到它。你知道这种行为是从哪里来的吗？这引起了一个问题，因为我们在这个实例上是多个用户，即使没有人使用它，GPU0也会饱和。

< P>如果您只对GPU1感兴趣，我会考虑将脚本包装成将代码> CUADAVISILBYX设备（参见）到<代码> 1 。这样，脚本只能看到一个GPU（它的id看起来像是

0

）。如果您将其设置为

2,3

，您将分别获得ID为

0,1

的GPU。

我们发现的一个解决方法是修改

skflow.TensorFlowEstimator

罪魁祸首是

with self._graph.as_default():
    tf.set_random_seed(self.tf_random_seed)
    self._global_step = tf.Variable(
        0, name="global_step", trainable=False)

在

skflow.TensorFlowEstimator.setup\u training（）

中，我们将其修改为

with self._graph.as_default(), tf.device("/gpu:{0}".format(self.gpu_number)):
    tf.set_random_seed(self.tf_random_seed)
    self._global_step = tf.get_variable('global_step', [],
                                      initializer=tf.constant_initializer(0), trainable=False)

---

在类中添加属性

gpu\U编号

，并在

skflow.TensorFlowEstimator中使用
allow\u soft\u placement=True
初始化
会话
。\u setup\u training（）

谢谢您的回答。你想提交一份PR，让它更通用，这样其他人就可以更容易地使用它吗？@YuanTang我一点都不熟悉pull请求，但我会在这个周末尝试一下。太好了！只需分叉存储库，在本地克隆分叉存储库，然后处理分叉回购。完成修改后，当您转到tensorflow/skflow中的pull request部分时，单击

跨叉比较

。如果您有任何问题，请告诉我，我很乐意为您提供帮助。