当前位置: 首页>>技术问答>>正文


python - Tensorflow在jupyter中设置CUDA_VISIBLE_DEVICES

技术问答 , , , 去评论

问题描述

我有两个GPU,想同时通过2个ipynb运行两个不同的网络,但是第一个notebook始终分配&占用了两个GPU。

使用CUDA_VISIBLE_DEVICES,我可以在python文件中隐藏设备,但是我不确定如何在jupyter notebook中执行此操作。

那么,能否将某个GPU设置为在同一服务器上运行的某个notebook不可见?或者说设置GPU对其他notebook隐藏?

 

最佳回答

您可以使用os.environ在笔记本中设置环境变量。在初始化TensorFlow以便将TensorFlow限制为第一个GPU之前,请执行以下操作。

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

您可以再次检查您是否拥有对TF可见的正确设备

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

我倾向于从notebook_util之类的实用程序模块中使用它

import notebook_util
notebook_util.pick_gpu_lowest_memory()
import tensorflow as tf

 

次佳回答

您可以通过使用magics更快地执行此操作,而无需任何导入:

%env CUDA_DEVICE_ORDER=PCI_BUS_ID
%env CUDA_VISIBLE_DEVICES=0

注意,所有env变量都是字符串,因此无需使用引号""。您可以通过运行%env <name_of_var>来验证是否已设置环境变量。或使用%env检查所有这些信息。

GPU Tensorflow
GPU设备 Tensorflow

参考资料

 

本文由《纯净天空》出品。文章地址: https://vimsky.com/article/4302.html,未经允许,请勿转载。