pytorch 加载大数据集 内存不够 的处理方式

文章用于记录一下本人对于较大数据集加载的问题的一些解决办法和思考。（比较口水话）

本文解决的问题适用领域：

1. 数据集很大，无法一次加载到内存
2. 纯文本类型的数据

我之前训练某个特定任务，习惯于把数据提前预处理为dataset保存起来【参见这里】，然后每次训练的时候直接加载这个文件。我这样做的目的是，方便调代码，使用很小量的数据先把代码调通，当出现一些小问题时不至于数据处理很久。

但是最近在使用自己的一些语料微调一个语言模型, 语料纯文本大小大概在8G左右，使用预处理脚本CPU满速处理后使用torch.save()保存成二进制文件，大概也8G左右。

这个时候问题来了，我使用了8张卡跑代码，这样的话就是8个进程，每个进程都要torch.load()加载这个二进制文件，256G的内存直接爆掉了，留下我一脸懵逼。

好，遇到了问题，开始动手解决。

内存不足？那多少够用呢？我先使用python窗口，单进程加载了这个文件，内存使用了60+G，算下来的话8卡是要500G的内存起步，不太现实，那么就从代码和数据的角度来解决吧。

使用原代码的处理方法？他处理后的存储sample应该也会占用60G内存，并且前期特别耗时，大概要四五小时，一旦发现后边的代码出了问题，相当于白处理了。

然后百度查相关的文档，看看有没有大佬解决这个问题。

有人说把数据分成若干份，每加载一个文件train起来，但是作者说这样做，内存依然会一直上升。有人说分块读取，但是这么大的数据量一直放在内存里，占用应该和前边一样也是特别大的，不一定够。

想来想去还是按照torch提供的dataset来处理，其实这是使用最广泛的方法，在getitem方法里写上处理逻辑就行了，返回相应的一条数据（sample）。但是现在的问题是，我读取了10行纯文本，每一行可能都会处理成若干条sa