Tensorflow set CUDA_VISIBLE_DEVICES within jupyter

后端 未结 2 1165
终归单人心
终归单人心 2020-11-28 18:19

I have two GPUs and would like to run two different networks via ipynb simultaneously, however the first notebook always allocates both GPUs.

Using CUDA_VISIBLE_DEV

相关标签:
2条回答
  • 2020-11-28 19:01

    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.

    0 讨论(0)
  • 2020-11-28 19:02

    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
    
    0 讨论(0)
提交回复
热议问题