完全理解你的痛点——训练大模型时GPU内存不够，纯CPU又慢到没法接受，还好TensorFlow确实有办法让你在GPU和主存之间做内存交换，同时保留大部分GPU加速的优势，下面是几个实用的方案：

1. 先启用GPU内存增长模式（基础配置） #

首先建议开启内存增长，让TensorFlow按需分配GPU内存，而不是一开始就占用全部GPU显存，这能避免不必要的内存浪费。代码很简单：

import tensorflow as tf

# 获取GPU设备
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        # 对每个GPU启用内存增长
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        logical_gpus = tf.config.list_logical_devices('GPU')
        print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
    except RuntimeError as e:
        # 内存增长必须在初始化TensorFlow之前设置
        print(e)

这个步骤能让TensorFlow只占用当前需要的显存，不会一开始就把GPU显存占满，为后续的内存交换留出空间。

2. 启用软设备放置（核心：自动在GPU/CPU间交换张量） #

这就是你要的“GPU内存交换到主存”的核心功能——TensorFlow的软设备放置（Soft Device Placement）。当某个张量太大，GPU装不下时，TensorFlow会自动把它放到CPU内存里，等需要在GPU上计算时再把数据传过去，虽然会有一些数据传输的开销，但比纯CPU训练快得多。

开启方法也很简单，在代码开头添加：

tf.config.set_soft_device_placement(True)

这个设置会让TensorFlow自动处理设备分配，不用你手动指定每个张量的位置。比如当你定义模型时，原本要放在GPU的层如果显存不够，TensorFlow会自动把部分参数或中间张量放到CPU，训练时自动完成数据的来回传输。

3. 混合精度训练（减少内存占用，可能直接解决问题） #

如果开启软设备放置后还是觉得效率损失有点大，可以先试试混合精度训练——用16位浮点数存储部分张量，能直接把显存占用减少约50%，很多情况下甚至不需要内存交换就能装下11GB的模型。

开启混合精度的代码：

from tensorflow.keras import mixed_precision

mixed_precision.set_global_policy('mixed_float16')

这个方法几乎不会损失模型精度（TensorFlow会自动在关键步骤用32位计算保证精度），还能加快训练速度，非常推荐优先尝试。

4. 模型并行（针对超大模型的进阶方案） #

如果你的模型实在太大，上面的方法还不够，可以考虑模型并行——把模型的不同层分配到GPU和CPU上，比如把一部分卷积层放GPU，把全连接层放CPU，手动控制张量的位置。示例代码大概是这样：

# 定义模型时手动指定设备
with tf.device('/GPU:0'):
    layer1 = tf.keras.layers.Dense(2048)
with tf.device('/CPU:0'):
    layer2 = tf.keras.layers.Dense(1024)

# 构建模型
input = tf.keras.Input(shape=(1000,))
x = layer1(input)
x = layer2(x)
output = tf.keras.layers.Dense(10)(x)
model = tf.keras.Model(inputs=input, outputs=output)

这个方法需要你手动规划模型的设备分配，适合对模型结构比较熟悉的场景，能更精准地控制内存使用。

总结一下：优先开启内存增长+混合精度，大部分情况下能直接解决显存问题；如果还是不够，再开启软设备放置，让TensorFlow自动做内存交换；极端情况再考虑模型并行。这些方法都能让你避免纯CPU训练的慢速度，同时接受可承受的效率损失。

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