當前位置: 首頁>>技術問答>>正文


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/zh-tw/article/4302.html,未經允許,請勿轉載。