TensorFlow 2.16中set_session如何使用？导入失败求助

在TensorFlow 2.16+搭配Keras3的环境里，set_session这个API已经被彻底移除了——毕竟TF2默认用的是即时执行模式，不再依赖Session来管理计算图和资源，所以你之前试的各种导入方式都找不到这个函数。

不用纠结旧API，直接用TF2原生的GPU配置方案就能替代原来的内存占比需求，下面是具体实现：

1. 替代原get_session的固定内存占比配置 #

这个函数和你原来的逻辑一致，设置GPU内存的使用比例：

import tensorflow as tf

def configure_gpu_memory(fraction):
    gpus = tf.config.list_physical_devices('GPU')
    if gpus:
        try:
            for gpu in gpus:
                # 获取GPU总内存，计算要分配的额度
                total_memory = tf.config.experimental.get_memory_info(gpu)['total']
                memory_limit = int(total_memory * fraction)
                # 应用内存限制配置
                tf.config.set_logical_device_configuration(
                    gpu,
                    [tf.config.LogicalDeviceConfiguration(memory_limit=memory_limit)]
                )
            print(f"GPU内存占比已设为{fraction}")
        except RuntimeError as e:
            print(f"配置失败: {e}")

2. 更推荐的动态内存分配方式 #

TF2官方更推荐这种方式，让GPU根据实际需求动态分配内存，不会一开始就占满显存：

import tensorflow as tf

def enable_gpu_memory_growth():
    gpus = tf.config.list_physical_devices('GPU')
    if gpus:
        try:
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu, True)
            print("已启用GPU内存动态分配")
        except RuntimeError as e:
            print(f"启用失败: {e}")

使用方法 #

一定要在导入其他库、定义模型之前调用这些配置函数，比如：

# 项目代码最开头执行配置
configure_gpu_memory(0.6)  # 比如设置用60%的GPU内存
# 或者选择动态分配
# enable_gpu_memory_growth()

# 之后再写你的模型代码、训练逻辑...

如果你的项目里还有其他依赖TF1.x Session的旧代码，尽量逐步迁移到TF2的即时执行模式。实在要兼容的话，可以用tf.compat.v1.disable_eager_execution()强制关闭即时执行，但这不是长久之计，后续版本可能会彻底砍掉相关兼容接口。

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