如何有效解决TensorFlow与Keras处理大数据量时内存溢出难题？

本文共计1290个文字，预计阅读时间需要6分钟。

内存溢出问题是参与Kaggle比赛或进行大数据量实验时的常见拦路虎。先前的小项目练习使新手产生产生了一个习惯性思维——读取训练集图片时，一次性将所有图片读入内存，然后分批进行训练。这种做法虽方便，但易导致内存不足。

内存溢出问题是参加kaggle比赛或者做大数据量实验的第一个拦路虎。

以前做的练手小项目导致新手产生一个惯性思维——读取训练集图片的时候把所有图读到内存中，然后分批训练。

其实这是有问题的，很容易导致OOM。现在内存一般16G，而训练集图片通常是上万张，而且RGB图，还很大，VGG16的图片一般是224x224x3，上万张图片，16G内存根本不够用。这时候又会想起——设置batch，但是那个batch的输入参数却又是图片，它只是把传进去的图片分批送到显卡，而我OOM的地方恰是那个“传进去”的图片，怎么办？

解决思路其实说来也简单，打破思维定式就好了，不是把所有图片读到内存中，而是只把所有图片的路径一次性读到内存中。

大致的解决思路为：

将上万张图片的路径一次性读到内存中，自己实现一个分批读取函数，在该函数中根据自己的内存情况设置读取图片，只把这一批图片读入内存中，然后交给模型，模型再对这一批图片进行分批训练，因为内存一般大于等于显存，所以内存的批次大小和显存的批次大小通常不相同。

下面代码分别介绍Tensorflow和Keras分批将数据读到内存中的关键函数。Tensorflow对初学者不太友好，所以我个人现阶段更习惯用它的高层API Keras来做相关项目，下面的TF实现是之前不会用Keras分批读时候参考的一些列资料，在模型训练上仍使用Keras，只有分批读取用了TF的API。