InfiniBatch

Infinibatch是一个用于深度神经网络训练中大规模数据集随机加载数据的检查点迭代器库。

特性

• 支持比内存能容纳的更大的语料库
• 整个语料库的分层块+句子级随机化，每个epoch有不同的随机化
• 只加载所需的数据
• 非常快的启动时间（不需要读取整个语料库）
• 只需要最基本的数据准备（例如，无需索引）
• 对于多GPU，只加载各自的GPU需要的部分
• 100%准确的检查点，从检查点恢复时不应读取所有数据直到检查点
• 支持带有动态批量大小的自动分桶批处理
• 预取线程
• 可组合的，支持复杂批处理，例如从多个文档中获取负样本

入门指南

Infinibatch需要Python 3.6或更高版本，并且没有依赖项。目前没有pip包。

要安装它，请克隆此存储库并在本地安装。

git clone https://github.com/microsoft/infinibatch
cd infinibatch
pip install -e .

文档

文档可以在这里找到： https://microsoft.github.io/infinibatch/

教程

这个简短教程将指导您如何准备数据，并从Python代码中以批量的方式消费它们。

Infinibatch基础：迭代器和检查点

Infinibatch提供Python迭代器来读取您的数据。迭代器表示可以逐个检索数据的数据流，例如通过一个for循环或重复调用next()。

Infinibatch对项目的数据类型一无所知，这由用户提供的文件读取函数确定。在NLP应用中，项目通常是文本的元组。在其他应用中，它们可以是图像或带有文本注释的音频文件。

Infinibatch使您能够以随机顺序读取数据，并支持检查点，这使得您可以从头开始训练。

随机化是即时进行的，这意味着不需要将整个数据集读入内存进行洗牌。Infinibatch实现了一个分层洗牌算法，在任何时候只持有数据的一部分在RAM中。

Infinibatch迭代器是可检查点的。检查点允许您在任何时候检索数据流中的当前位置（即“检查点”），这样您就可以稍后“回滚”到该位置。不幸的现实是，长时间运行的训练有时会崩溃。为了能够像没有崩溃一样继续崩溃的训练，在训练过程中保存中间模型时，请将Infinibatch迭代器的检查点保存到磁盘。要重新启动崩溃的训练，请将迭代器重置为保存的检查点。数据读取器现在将产生与崩溃时相同的相同数据项序列。

数据准备

Infinibatch对您的数据组织有一个要求：要使用Infinibatch中的数据，它必须分成大量的小块。块是从磁盘加载到RAM中的最小数据单元。Infinibatch在内存中保留块的一个随机子集，从中随机抽取样本。

以下我们将展示如何创建这样的分割。将数据分割成块的一种简单方法是使用Linux split命令。

在这个教程中，我们的“语料库”由6行文本组成，其中每行是一个数据项。要创建这个语料库，请在bash shell中运行以下命令。它创建一个名为corpus.txt的6行文本文件

echo \\
'Lorem ipsum dolor sit amet,
consectetur adipiscing elit,
sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.
Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.
The quick brown fox jumps over the lazy dog.' \\
> corpus.txt

现在让我们将其分成3个2行的块。每个块都存储为压缩文本文件。我们将在名为corpus_chunks的新子目录中创建它们

mkdir corpus_chunks
split  --lines 2  --numeric-suffixes                 \\
       --filter 'gzip > corpus_chunks/$FILE.txt.gz'  \\
       corpus.txt  corpus.

这将创建三个文件：corpus_chunks/corpus.00.txt.gz、corpus_chunks/corpus.01.txt.gz和corpus_chunks/corpus.02.txt.gz。要验证数据是否按预期分割，您可以使用此命令

zcat corpus_chunks/corpus.*.txt.gz

提示：对于大型语料库，我们建议用pigz（apt-get install pigz）替换gzip，它通过多线程运行明显更快。

使用Infinibatch按随机顺序读取项

我们将首先展示使用辅助函数chunked_dataset_iterator``()读取数据的简单方法。此函数将创建一个Infinibatch迭代器，以随机顺序产生您数据的内容。请看以下程序

import gzip, glob

from infinibatch import datasets as ds

ds = ds.chunked_dataset_iterator(
    chunk_refs = glob.glob('corpus_chunks/corpus.*.txt.gz'),
    read_chunk_fn = lambda path: iter(gzip.decompress(open(path, "rb")  \\
                                      .read()).decode(encoding='utf-8') \\
                                      .splitlines()),
    buffer_size = 6, seed = 1)

for i in range(10):
    print(next(ds))

您应该得到包含6个示例行的随机化输出

Lorem ipsum dolor sit amet,
consectetur adipiscing elit,
Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.
Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.
The quick brown fox jumps over the lazy dog.
sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.
consectetur adipiscing elit,
Lorem ipsum dolor sit amet,
The quick brown fox jumps over the lazy dog.
sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.

注意：buffer_size参数确定在任何给定时间读入内存的句子数，以从中抽取随机项。在实际设置中，对于包含数亿行文本的语料库，应将buffer_size参数设置为百万级别。RAM使用量和启动时间将与缓冲区大小成正比（但远低于将整个语料库加载到RAM中）。

批量读取不同长度的项

对于深度学习，我们希望将多个项目分组到批次中。对于NLP任务，项目通常是不同长度的文本行。Infinibatch实现了一种算法，该算法随机化输入序列，并将其分组为长度大致相同的批次（也称为分桶）。

Infinibatch的BucketedReadaheadBatchIterator执行此任务。它实现了一种基于Marian工具包的算法，预加载大量随机项目（通常是数百万；在这个例子中：6），对它们进行排序，并将它们分组为类似长度的批次，然后按顺序随机提供。

以下是一个示例。请注意，BucketedReadaheadBatchIterator接受先前的随机句子序列迭代器（ds）作为随机项目的来源。这是一个示例，说明了如何使用Infinibatch形成迭代器管道（这是一个类似于Python自身itertools的概念）。一旦将迭代器传递给另一个迭代器作为其来源，请将其视为由另一个迭代器拥有的，它必须不再被调用代码访问。

import gzip, glob

from infinibatch import datasets as ds
from infinibatch import iterators as it

ds = ds.chunked_dataset_iterator(
    chunk_refs = glob.glob('corpus_chunks/corpus.*.txt.gz'),
    read_chunk_fn = lambda path: iter(gzip.decompress(open(path, "rb")  \\
                                      .read()).decode(encoding='utf-8') \\
                                      .splitlines()),
    buffer_size = 6, seed = 1)

bs = it.BucketedReadaheadBatchIterator(
    source_iterator = ds,   # note: this is the iterator from above
    read_ahead = 6,
    key = lambda line: len(line),
    batch_size = 2,
    seed = 1)

for i in range(25):
    print(next(bs))

此代码应输出类似以下内容

['sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.',
 'The quick brown fox jumps over the lazy dog.']
['consectetur adipiscing elit,', 'Lorem ipsum dolor sit amet,']
['Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.',
 'Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.']

然后是相同元组的不同排列。如您所见，句子是随机顺序的，并且以长度大致相同的2批分组。您可能会注意到，项目分组到批次中的方式没有变化--这是本例中的情况，但在实际使用中通常不是这样，当数据大小远大于批次大小时。

在NLP中，句子长度往往差异很大。因此，使用像上面示例中那样的固定行数的批次将浪费GPU RAM和核心。这是因为行数受最长可能序列的限制；较短的行批次会留下GPU周期未被利用。理想情况下，人们会使用包含尽可能多的行的批次，前提是给定批次中最长行的令牌数。为了支持可变批次大小，Infinibatch允许将函数传递给batch_size参数。该函数将获得批次的最长项目，并应估计最多可以容纳多少个长度不超过此长度的项目。

在我们的示例中，我们假设批次最多可以容纳150个令牌。请按照以下方式更改上述代码

    batch_size = lambda longest_line: 150 // len(longest_line),

输出看起来像这样

['consectetur adipiscing elit,', 'Lorem ipsum dolor sit amet,']
['Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.']
['sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.',
 'The quick brown fox jumps over the lazy dog.']
['Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.']

较短的句子被分组，而较长的句子没有，因为它们会超过150个字符的总数。

将批次读入NumPy数组

最后，我们需要将批次喂入我们最喜欢的深度学习工具。我们将展示如何将文本行批次转换为填充的numpy数组。

在典型的NLP应用中，文本项目会被标记化，然后每个标记都由单元词汇表中的一个索引表示。为了简单起见，在这个例子中，每个字符都是一个标记，每个标记的数值单元索引只是它的ASCII码。然后，这些序列被填充到相等的长度，并用-1替换，然后转换为一个numpy数组。

请重新运行前面的示例，但在最后的for循环之前插入以下代码。此示例使用Infinibatch的MapIterator，它将用户提供的函数或lambda应用于每个项目

import numpy as np
def collate(lines_batch):
    # tokenize all lines in the batch and map to unit ids
    ids_batch = [[ord(c) for c in line] for line in lines_batch]
    # create a padded numpy array as wide as the longest line,
    # where shorter sequences are padded with -1
    width = max(len(ids) for ids in ids_batch)
    return np.array([ids + [-1] * (width-len(ids)) for ids in ids_batch])

bs = it.MapIterator(
    source_iterator = bs,
    transform = collate)

这将输出类似以下内容的批次。请注意，在包含多个句子的批次中，一些条目被用-1填充。

[[ 99 111 110 115 101  99 116 101 116 117 114  32  97 100 105 112 105 115
   99 105 110 103  32 101 108 105 116  44]
 [ 76 111 114 101 109  32 105 112 115 117 109  32 100 111 108 111 114  32
  115 105 116  32  97 109 101 116  44  -1]]
[[ 85 116  32 101 110 105 109  32  97 100  32 109 105 110 105 109  32 118
  101 110 105  97 109  44  32 113 117 105 115  32 110 111 115 116 114 117
  100  32 101 120 101 114  99 105 116  97 116 105 111 110  32 117 108 108
   97 109  99 111  32 108  97  98 111 114 105 115  32 110 105 115 105  32
  117 116  32  97 108 105 113 117 105 112  32 101 120  32 101  97  32  99
  111 109 109 111 100 111  32  99 111 110 115 101 113 117  97 116  46]]
[[115 101 100  32 100 111  32 101 105 117 115 109 111 100  32 116 101 109
  112 111 114  32 105 110  99 105 100 105 100 117 110 116  32 117 116  32
  108  97  98 111 114 101  32 101 116  32 100 111 108 111 114 101  32 109
   97 103 110  97  32  97 108 105 113 117  97  46]
 [ 84 104 101  32 113 117 105  99 107  32  98 114 111 119 110  32 102 111
  120  32 106 117 109 112 115  32 111 118 101 114  32 116 104 101  32 108
   97 122 121  32 100 111 103  46  -1  -1  -1  -1  -1  -1  -1  -1  -1  -1
   -1  -1  -1  -1  -1  -1  -1  -1  -1  -1  -1  -1]]
[[ 68 117 105 115  32  97 117 116 101  32 105 114 117 114 101  32 100 111
  108 111 114  32 105 110  32 114 101 112 114 101 104 101 110 100 101 114
  105 116  32 105 110  32 118 111 108 117 112 116  97 116 101  32 118 101
  108 105 116  32 101 115 115 101  32  99 105 108 108 117 109  32 100 111
  108 111 114 101  32 101 117  32 102 117 103 105  97 116  32 110 117 108
  108  97  32 112  97 114 105  97 116 117 114  46]]

从这里如何继续

上面的教程向您展示了如何使用最常用的迭代器类型，这是通过便利函数chunked_dataset_iterator()创建的。

本函数并不涵盖所有实际场景。例如，多任务学习场景需要更复杂的数据组合。要创建这些组合，您需要从底层构建块中组合必要的数据读取器。这在iterators模块的文档中有详细说明。

文档

要查看文档，请克隆存储库并转到docs/infinibatch/index.html

在编写文档时，请安装pdoc

pip install pdoc3

然后您可以启动一个本地http服务器，该服务器可以动态更新文档

pdoc --template-dir docs --http : infinibatch

我们目前尚未设置CI来自动生成文档。在将任何内容合并到master之前，请删除docs/infinibatch中的现有文档并运行

pdoc -o docs --template-dir docs --html infinibatch

测试

要运行单元测试，请运行以下命令。

python -m unittest discover -s test

如果您希望单元测试在第一个失败的测试后停止，请使用

python -m unittest discover -s test --failfast

使用mypy进行类型检查（如果已安装）

mypy infinibatch

贡献

本项目欢迎贡献和建议。大多数贡献都需要您同意一份贡献者许可协议（CLA），声明您有权，并且实际上确实授予我们使用您贡献的权利。有关详细信息，请访问https://cla.opensource.microsoft.com。

当您提交拉取请求时，CLA机器人将自动确定您是否需要提供CLA，并相应地装饰PR（例如，状态检查，注释）。只需遵循机器人提供的说明即可。您只需在整个使用我们CLA的所有存储库中这样做一次。

本项目已采用微软开源行为准则。有关更多信息，请参阅行为准则常见问题解答或通过opencode@microsoft.com联系以提出任何额外的问题或评论。