TensorFlow模型训练迭代内存持续增长问题求助（附复现代码）

你碰到的每轮迭代内存持续上涨的情况，在TensorFlow 2.0这个早期版本里属于比较典型的问题，结合你的代码，我们可以从几个关键点入手解决：

一、先修正代码里的明显错误 #

看你的create_model函数，循环添加Dense层时犯了个小错误：每一层都指定了input_shape。在Sequential模型里，只有第一层需要设置input_shape来定义输入维度，后面的层会自动根据前一层的输出形状推断输入，重复设置只会生成多余的张量节点，长期循环下来会累积内存占用。

修正后的代码应该是这样：

# 记得补充导入Sequential
from tensorflow.keras.models import Sequential

def create_model(self, dense_params=[256]):
    model = Sequential()
    # 仅第一层设置input_shape
    model.add(Dense(dense_params[0], activation='relu', input_shape=[self.state_size]))
    # 后续层直接添加即可
    for params in dense_params[1:]:
        model.add(Dense(params, activation='relu'))
    model.add(Dense(self.action_size, activation="linear"))
    model.compile(loss="mse", optimizer=Adam(lr=self.learning_rate))
    return model

二、循环调用model.fit是内存泄漏的核心原因 #

TensorFlow 2.0的早期版本中，反复在循环里用model.fit处理单个样本，会触发很多不必要的内部操作：每次fit都会临时构建计算图、初始化日志缓存，这些资源在循环中无法及时释放，就会导致内存越用越多。

这里给你两个更高效的替代方案：

方案1：用train_on_batch替代fit #

train_on_batch是专门的单步训练API，开销比fit小很多，适合循环里的单样本/小批量训练：

for i in range(10_000):
    state = np.random.rand(Agent.state_size)
    state = np.expand_dims(state, axis=0)
    output = np.random.rand(Agent.action_size)
    output = np.expand_dims(output, axis=0)
    # 用train_on_batch执行单步训练
    loss = Agent.model.train_on_batch(state, output)
    # 每隔一定步数打印日志，避免频繁输出的缓存占用
    if i % 100 == 0:
        print(f"第{i}轮迭代，当前损失：{loss}")

方案2：用tf.data.Dataset批量处理数据 #

如果你的训练数据可以提前生成或者流式读取，推荐用tf.data.Dataset来管理数据，一次性调用fit完成训练，这样能彻底避免循环调用带来的资源浪费：

# 批量生成训练数据
def generate_train_data(num_samples):
    states = np.random.rand(num_samples, Agent.state_size)
    outputs = np.random.rand(num_samples, Agent.action_size)
    return states, outputs

# 生成10000条样本
train_states, train_outputs = generate_train_data(10_000)
# 转换为TensorFlow数据集，按批次处理
train_dataset = tf.data.Dataset.from_tensor_slices((train_states, train_outputs)).batch(1)
# 一次性启动训练
Agent.model.fit(train_dataset, epochs=1, verbose=1)

三、升级TensorFlow版本（最直接的解决方式） #

TensorFlow 2.0作为2.x系列的第一个正式版，存在不少已知的内存泄漏bug，后续的2.2+版本已经修复了大量这类问题。如果你的项目环境允许，建议升级到较新的稳定版本，这能从根源上解决很多类似的内存问题。

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