使用Keras Sequential模型predict时遇Tensorflow显存不足错误求助

1. 分批预测并即时处理结果（避免累积所有预测张量） #

如果之前拆分测试集后仍在累积所有预测结果，会导致GPU/主机内存被持续占用。改为每批预测后直接处理结果（如写入文件、计算指标），不保存完整的预测集合：

batch_size = 16  # 进一步缩小批次规模
for idx in range(0, len(test_data), batch_size):
    batch_data = test_data[idx:idx+batch_size]
    batch_preds = model.predict(batch_data, verbose=0)
    # 即时处理当前批次结果，例如写入本地文件
    # with open("predictions.txt", "a") as f:
    #     for pred in batch_preds:
    #         f.write(f"{pred}\n")
    # 不要将batch_preds加入全局列表，避免内存堆积

2. 切换到CPU执行预测 #

你的主机内存有64GB，完全足够容纳大张量。强制将模型和测试数据转移到CPU运行：

import tensorflow as tf

# 将模型切换到CPU
model = model.to("cpu")

# 若测试数据是tf.data.Dataset，确保其在CPU上运行
test_data_cpu = test_data.map(
    lambda x, y: (tf.convert_to_tensor(x, dtype=tf.float32), y),
    num_parallel_calls=tf.data.AUTOTUNE
).prefetch(tf.data.AUTOTUNE)

# 执行预测
predictions = model.predict(test_data_cpu)

3. 检查测试数据的序列长度异常 #

该语言的测试集可能存在超长序列，导致模型输出的张量维度远超其他语言。先统计序列长度：

sequence_lengths = []
for x, _ in test_data:
    sequence_lengths.append(x.shape[0])

print(f"最长序列长度: {max(sequence_lengths)}")
print(f"平均序列长度: {sum(sequence_lengths)/len(sequence_lengths)}")

如果存在超长序列，对其单独截断或处理，避免影响整个批次的输出规模：

# 截断超长序列到固定长度（比如其他语言的最大序列长度）
max_seq_len = 256  # 替换为其他语言的最大序列长度
test_data_truncated = test_data.map(
    lambda x, y: (tf.keras.preprocessing.sequence.pad_sequences(x, maxlen=max_seq_len, truncating='post'), y)
)

4. 清理显存并关闭XLA优化 #

TensorFlow的XLA优化可能导致显存占用异常，同时手动清理显存后重新加载模型：

import tensorflow as tf

# 清理显存
tf.keras.backend.clear_session()
# 关闭XLA优化
tf.config.optimizer.set_jit(False)

# 重新加载模型（如果是从文件保存的模型）
model = tf.keras.models.load_model("your_trained_model.h5")
# 再执行预测
predictions = model.predict(test_data)

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