Cleverhans, как выбрать одно устройство с графическим процессором?

Question

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

Я хотел бы, однако, иметь возможность указать устройство GPU, которое я хотел бы использовать для обучения.

Я знаю о devices аргумент train функция, однако я пробовал несколько значений для этого поля, и оно всегда выделяет память для всех графических процессоров.

   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 исправления ошибки.