How to set specific gpu in tensorflow?

后端 未结 5 979
没有蜡笔的小新
没有蜡笔的小新 2020-11-29 21:45

I want to specify the gpu to run my process. And I set it as follows:

import tensorflow as tf
with tf.device(\'/gpu:0\'):
    a = tf.constant(3.0)
with tf.Se         


        
5条回答
  •  甜味超标
    2020-11-29 22:10

    There are 3 ways to achieve this:

    1. Using CUDA_VISIBLE_DEVICES environment variable. by setting environment variable CUDA_VISIBLE_DEVICES="1" makes only device 1 visible and by setting CUDA_VISIBLE_DEVICES="0,1" makes devices 0 and 1 visible. You can do this in python by having a line os.environ["CUDA_VISIBLE_DEVICES"]="0,1" after importing os package.

    2. Using with tf.device('/gpu:2') and creating the graph. Then it will use GPU device 2 to run.

    3. Using config = tf.ConfigProto(device_count = {'GPU': 1}) and then sess = tf.Session(config=config). This will use GPU device 1.

提交回复
热议问题