Cleverhans, как выбрать одно устройство GPU?

Question

Я тиражирую один из базовых уроков MNIST, используя cleverhans.У меня есть доступ к машине с несколькими графическими процессорами, и кажется, что библиотека использует все преимущества архитектуры мульти-графических процессоров во время обучения, и это здорово.

Я бы хотел, однако, иметь возможность указатьУстройство 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 должен быть доступ, состоит в добавлении CUDA_VISIBLE_DEVICES=0,1,2 к вашей команде python. Это будет использовать только графические процессоры 0, 1 и 2.

CUDA_VISIBLE_DEVICES=0,1,2 python script.py

Если в нашем методе train есть ошибка с аргументом устройств, не стесняйтесь открыть проблему или PR, исправляя ошибку.