Cleverhans，如何选择单个GPU设备？

Question

我正在使用cleverhans复制一个基本的MNIST教程。我可以使用多gpu机器，而且这个库似乎在训练过程中充分利用了多gpu架构，这很棒。

然而，我希望能够指定我想要用于培训的gpu设备。

我知道train函数的devices参数，但是我尝试了该字段的多个值，它总是在所有gpus上分配内存。

   train(
       sess,
       loss,
       x_train,
       y_train,
       devices=['/device:GPU:2', ],
       # also tried:
       # devices=["/GPU:0"],
       # devices=[2, ],
       # devices=['/gpu:2']
       # devices=['gpu:2']
       evaluate=evaluate,
       args=train_params,
       rng=rng
   )

有没有办法使用一个(或几个)特定的gpu，并且只在那里分配内存？

谢谢

Answer 1

指定希望python进程能够访问的GPU的另一种方法是在python命令中添加CUDA_VISIBLE_DEVICES=0,1,2。这将仅使用GPU 0、1和2。

CUDA_VISIBLE_DEVICES=0,1,2 python script.py

如果我们的train方法中的设备参数有一个bug，请随时打开一个问题或者PR修复这个bug。