TensorFlow跨脚本实例化时指定GPU失效问题求助

我之前也碰到过一模一样的问题，核心原因是TensorFlow会在首次导入或初始化阶段缓存GPU设备列表——你在FCN类的__init__里设置os.environ['CUDA_VISIBLE_DEVICES']的时机太晚了：当从其他脚本导入FCN类时，主脚本可能已经提前导入了TensorFlow或其他依赖TF的模块，这时TF已经完成了GPU检测，后续修改环境变量根本不会生效。

下面是两种可行的解决方案，按推荐度排序：

方法一：用TensorFlow原生API动态指定可见GPU（推荐） #

这种方法不依赖系统环境变量，能在运行时灵活控制GPU可见性，完全不受导入顺序影响。修改你的FCN.py中的__init__方法，把这段代码放在所有模型初始化逻辑的最前面：

import tensorflow as tf

class FCN:
    def __init__(self):
        # 获取系统中所有物理GPU设备
        physical_gpus = tf.config.list_physical_devices('GPU')
        if physical_gpus:
            try:
                # 仅让索引为1的GPU可见（对应你要指定的第2块GTX1080）
                tf.config.set_visible_devices(physical_gpus[1], 'GPU')
                # 验证配置：此时逻辑GPU数量应该为1
                logical_gpus = tf.config.list_logical_devices('GPU')
                print(f"已配置可见GPU：{len(logical_gpus)} 块（物理GPU总数：{len(physical_gpus)}）")
            except RuntimeError as e:
                # 提示错误：必须在创建TF会话/模型前执行配置
                print(f"GPU配置失败：{e}")
        
        # 此处编写你的模型初始化、权重加载等后续代码...

这样不管你从哪个脚本实例化FCN对象，TensorFlow都会严格只使用你指定的GPU，不会再抢占全部3块卡的资源。

方法二：提前设置环境变量 #

如果坚持用环境变量的方式，必须在主脚本的最开头、任何TensorFlow或FCN类导入之前设置，比如你的主脚本main.py：

import os
# 先设置环境变量，再导入任何依赖TF的模块
os.environ['CUDA_VISIBLE_DEVICES'] = '1'

# 之后再导入FCN类和其他模块
from FCN import FCN

# 实例化模型
fcn_model = FCN()

这种方法的局限性很明显：如果主脚本中其他模块在导入FCN前已经悄悄导入了TensorFlow，环境变量的设置还是会失效，所以可靠性不如方法一。

额外注意事项 #

如果你的应用中还有非TensorFlow的GPU模型（比如PyTorch），这些模型也可能抢占GPU资源，需要给它们单独配置：

• 对于PyTorch，可以用torch.cuda.set_device(1)，或者同样在主脚本开头设置os.environ['CUDA_VISIBLE_DEVICES'] = '1'（这个环境变量对多数GPU框架都生效）。

修改后验证效果：跨脚本实例化时，日志应该会和FCN.py内部运行时一致，只显示指定的1块GPU被创建。

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