python-3.x - Cleverhans，如何选择单个 GPU 设备？

Question

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

但是，我希望能够指定我想用于培训的 gpu 设备。

我知道函数的devices参数train，但是我已经为该字段尝试了多个值，它总是在所有 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 的另一种方法是CUDA_VISIBLE_DEVICES=0,1,2在 python 命令之前添加。这将仅使用 GPU 0、1 和 2。

CUDA_VISIBLE_DEVICES=0,1,2 python script.py

如果我们的方法中的设备参数存在错误train，请随时打开问题或 PR 修复错误。