如何修改代码避免Autokeras完成Trial 1后耗尽GPU内存？

我来帮你搞定这个问题！你遇到的是AutoKeras在连续搜索模型时的GPU内存占用不释放/泄漏问题，结合你的代码和场景，咱们可以从这几个方面调整：

1. 替换旧版TensorFlow GPU配置为TF2.x原生方案 #

你之前用的tf.compat.v1.Session是TF1.x的写法，在TF2.x环境下会和原生执行模式冲突，导致显存无法正常释放。换成TF2.x的显存管理方式，让TensorFlow按需分配显存，而不是一开始就占满所有显存。

2. 调小训练批次并添加资源清理步骤 #

AutoKeras默认的batch_size可能偏大，加上每轮Trial后没有主动清理资源，导致显存越积越多。咱们手动添加清理逻辑，并调小batch_size降低单轮显存占用。

修改后的完整代码 #

import numpy as np
import tensorflow as tf
from tensorflow.keras.datasets import mnist
import autokeras as ak
import gc

# TF2.x 原生GPU内存配置：按需分配，避免一次性占满显存
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        logical_gpus = tf.config.list_logical_devices('GPU')
        print(f"{len(gpus)} 物理GPU, {len(logical_gpus)} 逻辑GPU")
    except RuntimeError as e:
        # 必须在任何GPU操作前设置，所以放最前面
        print(e)

# 加载MNIST数据集
(x_train, y_train), (x_test, y_test) = mnist.load_data()

# 初始化ImageClassifier，调小batch_size减少显存占用
clf = ak.ImageClassifier(
    overwrite=True,
    max_trials=10,
    batch_size=32  # 从默认的64降到32，显存压力会小很多
)

# 训练（不用手动包裹tf.device，TF2.x会自动管理GPU）
clf.fit(x_train, y_train, epochs=2)

# 每轮Trial结束后手动清理资源，防止内存泄漏
gc.collect()
tf.keras.backend.clear_session()

额外优化方案（如果显存还是不够） #

如果你的GPU显存特别小（比如只有4G），可以直接限制显存使用量，比如只分配2GB给模型：

if gpus:
    try:
        tf.config.set_logical_device_configuration(
            gpus[0],
            [tf.config.LogicalDeviceConfiguration(memory_limit=2048)]
        )
        logical_gpus = tf.config.list_logical_devices('GPU')
        print(f"{len(gpus)} 物理GPU, {len(logical_gpus)} 逻辑GPU")
    except RuntimeError as e:
        print(e)

关键调整说明 #

• 移除手动tf.device包裹：TF2.x和AutoKeras会自动检测并使用GPU，手动指定设备反而可能导致资源释放不及时，交给框架自动管理更可靠。
• 显存增长模式：set_memory_growth让TensorFlow根据模型需求动态分配显存，训练结束后会自动释放不用的部分。
• 资源清理：gc.collect()回收Python层面的内存，tf.keras.backend.clear_session()清除Keras会话和模型占用的GPU资源，彻底避免内存泄漏。

这样调整后，每轮Trial结束后GPU内存应该会正常释放，不会出现Trial2启动时显存占满的情况了。

内容的提问来源于stack exchange，提问作者Петр Воротинцев