日常使用过程中往往需要涉及 GPU 进行模型训练和推理，及指定 GPU进行计算，那么：



经常地，我会这么使用定义要使用的 gpu_id 和 需要消耗的显存:

import tensorflow as tf
import numpy as np
from tensorflow import keras
print(tf.__version__)
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1'
# 获取物理gpu, cpu对象
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
cpus = tf.config.experimental.list_physical_devices(device_type='CPU')

# 设置当前程序的　物理可见设备范围
tf.config.experimental.set_visible_devices(devices=gpus[0], device_type='GPU')

#设置仅在需要时申请：
# for gpu in gpus:
#     tf.config.experimental.set_memory_growth(gpu, True)

# 设置在物理gpu上设置虚拟gpu，并用来限制gpu内存使用
tf.config.experimental.set_virtual_device_configuration(
    gpus[0], [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024), tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])

# 获取当前虚拟gpu对象
logical_gpus = tf.config.experimental.list_logical_devices('GPU')

print('物理gpu个数: ', len(gpus))
print('总的gpu个数，其中包括括虚拟Logical gpus: ', len(logical_gpus))