TensorFlow训练GPU显存不足(OOM)问题求助

你的问题是GPU显存（约1.6GB）不足以容纳当前batch size下的模型计算张量，并非你主动开启了显存限制，完全可以通过调整训练策略让GPU正常工作，不用被迫只用CPU。

以下是具体的优化方案：

1. 减小Batch Size #

这是最直接有效的方法，降低单次训练的样本数量，减少显存占用：

history = model.fit(training_set, validation_data=test_set, epochs=20, batch_size=8)

可以从16开始尝试，逐步降到能稳定运行的数值，训练速度仍会远快于CPU。

2. 启用显存动态分配 #

让TensorFlow按需分配显存，避免启动时就占用全部GPU显存，在模型定义前添加：

import tensorflow as tf
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)

3. 缩小输入图像尺寸 #

错误中显示的张量尺寸[64,64,224,224]对应224x224的输入图像，适当缩小尺寸（比如192x192、160x160）能显著降低显存占用，对模型精度的影响通常很小。

4. 开启混合精度训练 #

利用半精度（float16）存储部分张量，在不损失精度的前提下减少显存使用，RTX 3050 Ti完全支持该特性：

from tensorflow.keras.mixed_precision import set_global_policy
set_global_policy('mixed_float16')

注意要在模型构建和编译前执行这段代码。

5. 简化模型结构 #

如果模型是自定义的，可以减少卷积层的滤波器数量（比如把64个改为32个），或者移除冗余的全连接层，降低模型整体的显存需求。

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