Google Colab中运行Keras简单CNN模型调用model.fit时内存耗尽问题求助

我之前在Colab处理大尺寸嵌入数据时也碰到过一模一样的内存崩溃问题，咱们一步步拆解解决：

1. 先把嵌套列表转成连续内存的numpy数组 #

你用的Python嵌套列表在内存中的开销远大于连续存储的numpy数组，TensorFlow处理时还会额外占用内存来转换格式。先把数据转成真正的np.float16连续数组：

import numpy as np

# 转换为numpy数组并确保连续存储
x_train = np.array(x_train, dtype=np.float16)
x_train = np.ascontiguousarray(x_train)
x_test = np.array(x_test, dtype=np.float16)
x_test = np.ascontiguousarray(x_test)

这一步能直接减少内存的碎片化占用。

2. 大幅降低batch_size #

虽然32看起来不算大，但你的输入样本尺寸是400×300，32个样本的单batch张量加上模型中间层输出、优化器状态，很容易突破Colab的内存上限。先把batch_size降到8或者4试试，比如：

model.fit(x_train, y_train, batch_size=8, epochs=2, verbose=1, validation_data=(x_test, y_test))

等能正常运行后，再逐步调大到合适的数值。

3. 使用TensorFlow Dataset API分批加载数据 #

直接把整个大数组喂给model.fit会一次性将所有数据加载到内存，而Dataset API可以实现按需分批加载，还能异步预处理，大幅降低内存压力：

import tensorflow as tf

# 转换为Dataset并分批
train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
# 分批+预取，让训练和数据准备并行
train_ds = train_ds.batch(8).prefetch(tf.data.AUTOTUNE)

val_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))
val_ds = val_ds.batch(8)

# 用Dataset训练
model.fit(train_ds, epochs=2, validation_data=val_ds)

4. 训练前手动清理内存 #

Colab的后台可能残留之前运行的张量、模型占用内存，训练前手动清理一下：

import gc
import tensorflow as tf

gc.collect()
tf.keras.backend.clear_session()

这能释放被闲置的内存空间。

5. 开启混合精度训练 #

虽然你的输入已经是float16，但开启TensorFlow的混合精度可以让模型计算过程也用float16，同时保持权重的数值稳定性，进一步压缩内存占用：

from tensorflow.keras.mixed_precision import set_global_policy

# 设置混合精度策略
set_global_policy('mixed_float16')
# 注意：如果最后一层是sigmoid这类对数值精度敏感的激活，建议指定dtype为float32
# 比如把最后一层Dense改成：Dense(1, dtype='float32')

按这个顺序尝试，应该能解决你的内存耗尽问题——我当时就是靠调整数据格式和用Dataset API解决的类似情况。

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