将进程绑定到GPU逻辑设备：多进程场景下如何设置逻辑设备可见性

你提到的tf.config.set_visible_devices接口实际同时支持物理GPU、逻辑GPU两类设备对象，无需单独的逻辑设备可见性设置接口，按以下步骤操作即可实现需求：

前置要求 #

逻辑GPU拆分操作必须在所有TensorFlow上下文初始化、算子执行前完成，一旦GPU被占用就无法再修改配置。

具体实现步骤 #

以你提到的2物理GPU、拆分16个逻辑GPU、8个进程各分配2个逻辑GPU的场景为例：

步骤1：拆分物理GPU为逻辑GPU #

在所有进程启动前的初始化阶段（也可以在每个子进程刚启动的最开头执行，只要保证在TF初始化前即可）执行物理GPU拆分：

import tensorflow as tf

# 检测可用物理GPU
physical_gpus = tf.config.list_physical_devices('GPU')
if not physical_gpus:
    raise RuntimeError("未检测到可用GPU")

# 每个物理GPU拆分为8个逻辑GPU，总共2*8=16个
logical_gpu_per_physical = 8
for gpu in physical_gpus:
    tf.config.set_logical_device_configuration(
        gpu,
        # 按需求设置每个逻辑GPU的显存上限，单位为MB
        [tf.config.LogicalDeviceConfiguration(memory_limit=1024) for _ in range(logical_gpu_per_physical)]
    )

步骤2：子进程启动时分配专属可见逻辑GPU #

每个fit任务对应的子进程刚启动时，先获取全部逻辑GPU列表，再选择指定索引的逻辑GPU设置为可见，代码如下：

import tensorflow as tf

# 进程ID取值0~7，对应8个fit进程
process_id = 0
# 每个进程分配2个逻辑GPU
gpu_per_process = 2
# 计算当前进程分配的逻辑GPU索引范围
start_idx = process_id * gpu_per_process
end_idx = start_idx + gpu_per_process

# 获取全部逻辑GPU列表
all_logical_gpus = tf.config.list_logical_devices('GPU')
# 选择当前进程需要的逻辑GPU
target_gpus = all_logical_gpus[start_idx:end_idx]
# 设置仅指定逻辑GPU可见
tf.config.set_visible_devices(target_gpus, 'GPU')

# 验证配置：打印当前进程可见的GPU，确认只有2个逻辑GPU
print(tf.config.get_visible_devices('GPU'))

# 后续再执行fit等TensorFlow操作

注意事项 #

• 所有GPU配置操作必须放在TensorFlow初始化的最前端，一旦执行过任何TF算子、创建过张量、加载过模型，配置会直接失效
• 逻辑GPU的排序规则为：先按物理GPU的编号从小到大排列，同一块物理GPU下的逻辑GPU按创建顺序排列，上述示例中编号0_{7的逻辑GPU属于物理GPU0，8}15属于物理GPU1
• 如需跨物理GPU分配逻辑GPU，调整索引选择逻辑即可，例如让进程0拿索引0和8的逻辑GPU，分别对应两个物理GPU下的资源，实现跨卡负载均衡

内容的提问来源于stack exchange，提问作者Phylosopher