Cleverhans，如何select单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 方法中的设备参数存在错误，请随时打开问题或 PR 修复错误。

Cleverhans，如何select单GPU设备？

Cleverhans, how to select single GPU device?

python-3.x

tensorflow

cleverhans