使用Keras生成器调用model.fit时内存占用过高是否正常？

复现版本 #

在所有测试过的TensorFlow版本（2.7.0、2.7.3、2.8.0、2.9.0）中均复现了该异常，所有测试版本下都未能解决该问题。

运行环境配置 #

• 操作系统：Ubuntu 20
• GPU：RTX 2060
• 运行内存：16GB

自定义数据生成器实现 #

通过继承tf.keras.utils.Sequence实现自定义数据生成器向模型喂入数据，实现代码如下：

class DataGen(tf.keras.utils.Sequence):
    def __init__(self, indices, batch_size):
        self.X = X
        self.y = y
        self.indices = indices
        self.batch_size = batch_size
    
    def __getitem__(self, index):
        X_batch = self.X[self.indices][
            index * self.batch_size : (index + 1) * self.batch_size
        ]
        y_batch = self.y[self.indices][
            index * self.batch_size : (index + 1) * self.batch_size
        ]
        return X_batch, y_batch
    
    def __len__(self):
        return len(self.y[self.indices]) // self.batch_size

train_gen = DataGen(train_indices, 32)
val_gen = DataGen(val_indices, 32)
test_gen = DataGen(test_indices, 32)

上述代码中X、y为通过h5py从.h5文件加载的数据集，train_indices、val_indices、test_indices分别为训练集、验证集、测试集对应X、y的样本索引。

模型构建与训练代码 #

构建MobileNetV2模型并开展训练的代码如下：

# setup model
base_model = tf.keras.applications.MobileNetV2(input_shape=(128, 128, 3),
                                                include_top=False)
base_model.trainable = False

mobilenet1 = Sequential([
    base_model,
    Flatten(),
    Dense(27, activation='softmax')
])

mobilenet1.compile(optimizer=tf.keras.optimizers.Adam(),
                   loss=tf.keras.losses.CategoricalCrossentropy(),
                   metrics=['accuracy'])
# model training
hist_mobilenet = mobilenet1.fit(train_gen, validation_data=val_gen, epochs=1)

异常现象 #

训练启动前系统内存占用仅为8%，但训练一旦启动内存占用就从30%持续攀升至60%。使用生成器每次仅加载32条样本的小批量数据，如此高的内存涨幅十分异常。此外训练结束后内存占用仍维持在30%以上，检查所有全局变量，均无占用如此大内存的对象。如果再次启动训练，内存占用会进一步升高，最终导致Jupyter Notebook内核崩溃。

补充观测信息 #

1. 训练停止后内存占用会小幅下降，调用垃圾回收器可进一步降低内存占用，但即便删除fit方法返回的history对象，内存占用也无法回落至初始的8%水平。
2. 单个批次的X、y数据大小总和仅为48字节：每次仅加载48字节的数据却引发数GB级的内存涨幅完全不符合预期。使用HDF5数据集的初衷就是避免占用过高运行内存，曾猜测是fit方法内部创建了相关变量，但数GB的内存占用显然不符合逻辑。

核心疑问：该现象是代码实现存在问题，还是属于框架的正常表现？

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