张量流在木星内设置 CUDA_VISIBLE_DevICES

我有两个 GPU，并希望通过 ipynb 同时运行两个不同的网络，但是第一个笔记本总是分配两个 GPU。

使用 CUDA _ VISIBLE _ DISICES，我可以隐藏设备的巨蟒文件，但我不确定如何做到这一点在笔记本电脑。

是否有办法将不同的 GPU 隐藏到运行在同一台服务器上的笔记本电脑中？

224460 次浏览
小开
最佳答案

You can set environment variables in the notebook using os.environ. Do the following before initializing TensorFlow to limit TensorFlow to first GPU.

import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"   # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"]="0"

You can double check that you have the correct devices visible to TF

from tensorflow.python.client import device_lib
print device_lib.list_local_devices()

I tend to use it from utility module like notebook_util

import notebook_util
notebook_util.pick_gpu_lowest_memory()
import tensorflow as tf
小开

You can do it faster without any imports just by using magics:

%env CUDA_DEVICE_ORDER=PCI_BUS_ID
%env CUDA_VISIBLE_DEVICES=0

Notice that all env variable are strings, so no need to use ". You can verify that env-variable is set up by running: %env <name_of_var>. Or check all of them with %env.

小开

You can also enable multiple GPU cores, like so:

import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="0,2,3,4"

提交答案