有3种方法可以实现这一点：
1.使用CUDA_VISIBLE_DEVICES环境变量。通过设置环境变量CUDA_VISIBLE_DEVICES="1"，只显示设备1，通过设置CUDA_VISIBLE_DEVICES="0,1"，显示设备0和1。在python中，导入os包后，可以使用os.environ["CUDA_VISIBLE_DEVICES"]="0,1"行来完成此操作。
1.使用with tf.device('/gpu:2')并创建图形。然后它将使用GPU设备2运行。
1.使用config = tf.ConfigProto(device_count = {'GPU': 1})，然后使用sess = tf.Session(config=config)。这将使用GPU设备1。

TF会在每个可见的GPU上分配所有可用的内存，如果不告诉你的话。这里有5种方法可以坚持只使用一个（或几个）GPU。

• • Bash解决方案。**在启动python或jupyter notebook之前，在终端/控制台中设置CUDA_VISIBLE_DEVICES=0,1：

CUDA_VISIBLE_DEVICES=0,1 python script.py

• • Python解决方案。**在构造会话之前运行接下来的2行代码

import os
os.environ["CUDA_VISIBLE_DEVICES"]="0,1"

• • 自动解决方案。**以下方法将自动检测未被其他脚本使用的GPU设备，并为您设置CUDA_VISIBLE_DEVICES。您必须在构建会话之前调用mask_unused_gpus。它将根据当前内存使用情况过滤掉GPU。这样，您可以一次运行脚本的多个示例，而无需更改代码或设置控制台参数。

该功能：

import subprocess as sp
import os

def mask_unused_gpus(leave_unmasked=1):
  ACCEPTABLE_AVAILABLE_MEMORY = 1024
  COMMAND = "nvidia-smi --query-gpu=memory.free --format=csv"

  try:
    _output_to_list = lambda x: x.decode('ascii').split('\n')[:-1]
    memory_free_info = _output_to_list(sp.check_output(COMMAND.split()))[1:]
    memory_free_values = [int(x.split()[0]) for i, x in enumerate(memory_free_info)]
    available_gpus = [i for i, x in enumerate(memory_free_values) if x > ACCEPTABLE_AVAILABLE_MEMORY]

    if len(available_gpus) < leave_unmasked: raise ValueError('Found only %d usable GPUs in the system' % len(available_gpus))
    os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, available_gpus[:leave_unmasked]))
  except Exception as e:
    print('"nvidia-smi" is probably not installed. GPUs are not masked', e)

mask_unused_gpus(2)

局限性：如果你一次启动多个脚本可能会导致冲突，因为当你构造一个会话时内存不会立即分配。如果你遇到问题，你可以使用一个随机版本，就像在原始的source code: mask_busy_gpus()中一样

• • Tensorflow 2.0**建议使用另一种方法：

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  # Restrict TensorFlow to only use the first GPU
  try:
    tf.config.experimental.set_visible_devices(gpus[0], 'GPU')
  except RuntimeError as e:
    # Visible devices must be set at program startup
    print(e)

• • Tensorflow/Keras**还允许指定gpu与会话配置一起使用。只有在设置环境变量不是一个选项（即MPI运行）时，我才推荐使用它。因为它往往是所有方法中最不可靠的，尤其是与keras一起使用时。

config = tf.ConfigProto()
config.gpu_options.visible_device_list = "0,1"
with tf.Session(config) as sess:
#or K.set_session(tf.Session(config))

我相信你需要设置CUDA_VISIBLE_DEVICES=1。或者你想使用哪一个GPU。如果你只让一个GPU可见，不管你把环境变量设置成什么，你都会在张流中把它称为/gpu:0。
有关该环境变量的详细信息：https://devblogs.nvidia.com/cuda-pro-tip-control-gpu-visibility-cuda_visible_devices/

您可以通过在python脚本的开头添加以下内容来修改GPU选项设置：

gpu_options = tf.GPUOptions(visible_device_list="0")
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

“0”是你想要使用的GPU的名称。你可以在终端提示符中输入命令nvidia-smi来得到可用GPU的列表。
有了Keras，这2个功能允许选择CPU或GPU，在GPU的情况下，将使用内存的分数。

import os
from keras.backend.tensorflow_backend import set_session
import tensorflow as tf


def set_cpu_option():
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"  # see issue #152
    os.environ["CUDA_VISIBLE_DEVICES"] = ""
    os.environ["CUDA_VISIBLE_DEVICES"] = ""

def set_gpu_option(which_gpu, fraction_memory):
    config = tf.ConfigProto()
    config.gpu_options.per_process_gpu_memory_fraction = fraction_memory
    config.gpu_options.visible_device_list = which_gpu
    set_session(tf.Session(config=config))
    return

set_gpu_option("0", 0.9)
# or 
set_cpu_option()

def set_specific_gpu(ID):
    gpus_all_physical_list = tf.config.list_physical_devices(device_type='GPU')    
    tf.config.set_visible_devices(gpus_all_physical_list[ID], 'GPU')

参见https://www.tensorflow.org/api_docs/python/tf/config/list_physical_devices

import tensorflow as tf

gpu_number = 2 #### GPU number 
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    tf.config.experimental.set_visible_devices(gpus[gpu_number], 'GPU') 
    logical_gpus = tf.config.experimental.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPU")

在我的多核gpu设置中，我所见过的最优雅、最简洁的方式是：

import os
os.environ["CUDA_VISIBLE_DEVICES"]="1"
tf_device='/gpu:0'

这将任务分配给gpu设备1。
类似地，在线条上做某事：

import os 
os.environ["CUDA_VISIBLE_DEVICES"]="2"
tf_device='/gpu:0'

os.environ命令可以被看作是一种只暴露你想要运行代码的GPU设备的方法，第二个命令只选择你指定的第一个可用设备。

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

唯一的事情，我工作干净，从进程内分配特定的GPU到每个进程在一个池。

TF 2.9及更高版本的tensorflow已更改API，因此更新相同的API，

gpus = tf.config.list_physical_devices('GPU')
gpu_id = 0
if gpus:
    # Restrict TensorFlow to only use only one GPU based on gpu_id
    try:
        tf.config.set_visible_devices(gpus[gpu_id], 'GPU')
        logical_gpus = tf.config.list_logical_devices('GPU')
        print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPU")
    except RuntimeError as e:
        # Visible devices must be set before GPUs have been initialized
        print(e)