使用 XGBoost 外部内存版本

在大数据集上工作时，训练 XGBoost 模型可能会很具有挑战性，因为整个数据集需要加载到内存中。这可能会很昂贵，有时甚至不可行。从 1.5 版本开始，用户可以定义一个自定义迭代器，以分块方式加载数据来运行 XGBoost 算法。外部内存可以用于训练和预测，但训练是主要用例，这将是本教程的重点。对于预测和评估，用户可以自行遍历数据，而训练则需要将整个数据集加载到内存中。

在训练过程中，XGBoost 提供了两种不同的外部内存支持模式，一种用于基于 CPU 的算法，如 hist 和 approx，另一种用于基于 GPU 的训练算法。我们将在以下章节中介绍它们。

备注

exact 树方法不支持从外部内存数据进行训练。

备注

该功能在2.0版本中仍处于实验阶段。性能尚未得到良好优化。

外部内存支持经历了多次迭代，目前仍在积极开发中。与使用 DataIter 的 QuantileDMatrix 类似，XGBoost 使用用户提供的自定义迭代器逐批加载数据。然而，与 QuantileDMatrix 不同，外部内存不会在未使用 GPU 的情况下连接批次（它采用混合方法，详情如下）。相反，它会将所有批次缓存到外部内存中，并按需获取它们。请转到文档末尾查看 QuantileDMatrix 和外部内存之间的比较。

数据迭代器

从 XGBoost 1.5 开始，用户可以使用 Python 或 C 接口定义自己的数据加载器。demo 目录中有一些示例，供快速入门。这是文本输入外部内存的通用版本，用户不再需要准备 XGBoost 识别的文本文件。要启用该功能，用户需要定义一个包含 2 个类方法 next 和 reset 的数据迭代器，然后将其传递给 DMatrix 构造函数。

import os
from typing import List, Callable
import xgboost
from sklearn.datasets import load_svmlight_file

class Iterator(xgboost.DataIter):
  def __init__(self, svm_file_paths: List[str]):
    self._file_paths = svm_file_paths
    self._it = 0
    # XGBoost will generate some cache files under current directory with the prefix
    # "cache"
    super().__init__(cache_prefix=os.path.join(".", "cache"))

  def next(self, input_data: Callable):
    """Advance the iterator by 1 step and pass the data to XGBoost.  This function is
    called by XGBoost during the construction of ``DMatrix``

    """
    if self._it == len(self._file_paths):
      # return 0 to let XGBoost know this is the end of iteration
      return 0

    # input_data is a function passed in by XGBoost who has the exact same signature of
    # ``DMatrix``
    X, y = load_svmlight_file(self._file_paths[self._it])
    input_data(data=X, label=y)
    self._it += 1
    # Return 1 to let XGBoost know we haven't seen all the files yet.
    return 1

  def reset(self):
    """Reset the iterator to its beginning"""
    self._it = 0

it = Iterator(["file_0.svm", "file_1.svm", "file_2.svm"])
Xy = xgboost.DMatrix(it)

# The ``approx`` also work, but with low performance. GPU implementation is different from CPU.
# as noted in following sections.
booster = xgboost.train({"tree_method": "hist"}, Xy)

上面的代码片段是对外部内存的实验性支持的简化版本。对于C语言的示例，请参见 demo/c-api/external-memory/。迭代器是使用外部内存与XGBoost的通用接口，你可以将生成的 DMatrix 对象用于训练、预测和评估。

根据可用内存设置批量大小是很重要的。如果你有64GB的内存，一个好的起点是将批量大小设置为每批10GB。不建议设置像每批32个样本这样的小批量大小，因为这可能会严重损害梯度提升的性能。

CPU 版本

在上一节中，我们演示了如何使用 hist 树方法在 CPU 上训练基于树的模型。该方法涉及在树构建过程中迭代缓存中的数据批次。为了获得最佳性能，我们建议使用 grow_policy=depthwise 设置，这允许 XGBoost 仅通过几次批次迭代构建整个树层。相反，使用 lossguide 策略要求 XGBoost 为每个树节点迭代数据集，从而导致性能较慢。

如果使用外部内存，CPU训练的性能受限于IO（输入/输出）速度。这意味着磁盘IO速度主要决定了训练速度。在基准测试中，我们使用了连接到PCIe-4插槽的NVMe，其他类型的存储可能因速度过慢而无法实际使用。此外，您的系统可能会执行缓存以减少文件读取的开销。

GPU版本 (GPU Hist树方法)

外部内存由GPU算法支持（即当 device 设置为 cuda 时）。然而，用于GPU的算法与用于CPU的算法不同。在CPU上训练时，树方法在树构造算法的每个步骤中遍历外部内存中的所有批次。另一方面，GPU算法采用混合方法。它在每次迭代的开始遍历数据，并将所有批次连接到GPU内存中以提高性能。为了减少总体内存使用，用户可以利用子采样。GPU hist树方法支持基于梯度的采样，使用户能够在不牺牲准确性的情况下设置低采样率。

param = {
  ...
  'subsample': 0.2,
  'sampling_method': 'gradient_based',
}

关于采样算法及其在外存训练中的使用，请参阅这篇论文。

警告

当GPU在迭代外部内存时内存不足，用户可能会收到段错误而不是OOM异常。

备注

在使用外部内存与XGBoost时，数据被分割成更小的块，以便在任何给定时间只需将其中一部分存储在内存中。需要注意的是，这种方法仅适用于预测数据（X），而其他数据，如标签和内部运行时结构则会被连接在一起。这意味着，当处理宽数据集时，X``的大小明显大于其他数据如``y，内存减少的效果最为显著，而对于窄数据集，其影响则较小。

正如人们所料，按需获取数据会给存储设备带来巨大的压力。如今的计算设备可以处理的数据量远超存储设备在单位时间内能够读取的数据量。两者的比例相差几个数量级。一个GPU能够在瞬间处理数百GB的浮点数据。另一方面，连接到PCIe-4插槽的四通道NVMe存储通常具有约6GB/s的数据传输速率。因此，训练过程很可能会严重受限于存储设备。在采用外部内存解决方案之前，一些粗略的计算可能会帮助你判断其可行性。例如，如果你的NVMe驱动器每秒可以传输4GB（一个相当实际的数字）数据，并且你有一个100GB的压缩XGBoost缓存（这相当于一个大小约为200GB的密集float32 numpy数组，或多或少）。当参数正确时，深度为8的树至少需要16次数据遍历。在不考虑其他开销并假设计算与IO重叠的情况下，训练一棵树大约需要14分钟。如果你的数据集恰好有TB级别的大小，那么你可能需要数千棵树来获得一个泛化模型。这些计算可以帮助你估算预期的训练时间。

然而，有时我们可以改善这一限制。还应考虑到操作系统（主要讨论的是Linux内核）通常可以在主机内存中缓存数据。只有在有新数据进来且没有剩余空间时，才会驱逐页面。实际上，至少部分数据可以在整个训练过程中持久存在于主机内存中。我们在优化外部内存获取器时意识到了这种缓存。压缩缓存通常比原始输入数据小，特别是当输入是密集的且没有任何缺失值时。如果主机内存能够容纳这一压缩缓存的大部分，那么在初始化后性能应该是不错的。我们迄今为止的开发主要集中在两个方面的外部内存优化：

在适当的时候避免遍历数据。
如果操作系统能够缓存数据，性能应该接近于核心内训练。

从 XGBoost 2.0 开始，外部内存的实现使用了 mmap。它没有针对断开网络设备（如 SIGBUS）等系统错误进行测试。在面对总线错误时，您将看到硬崩溃，并需要清理缓存文件。如果训练会话可能需要很长时间，并且您正在使用 NVMe-oF 等解决方案，我们建议定期检查点保存模型。此外，值得注意的是，大多数测试都是在 Linux 发行版上进行的。

另一个需要记住的重要点是，为 XGBoost 创建初始缓存可能需要一些时间。与外部内存的接口是通过自定义迭代器，我们不能假设它们是线程安全的。因此，初始化是按顺序进行的。使用 xgboost.config_context 并设置 verbosity=2 可以在等待期间提供一些关于 XGBoost 正在做什么的信息，如果你不介意额外的输出。

与 QuantileDMatrix 相比

将迭代器传递给 QuantileDmatrix 可以直接用数据块构建 QuantileDmatrix。另一方面，如果将其传递给 DMatrix，它则启用了外部内存功能。QuantileDmatrix 在压缩后在内存中连接数据，并且在训练期间不获取数据。另一方面，外部内存 DMatrix 按需从外部内存获取数据批次。当你可以将大部分数据放入内存时，使用 `QuantileDMatrix`（如有必要使用迭代器）。训练速度将比使用外部内存快一个数量级。

文本文件输入

这是外部内存支持的原始形式，建议用户使用自定义数据迭代器。使用外部内存版本的文本输入与内存版本之间没有太大区别。唯一的区别是文件名格式。

外部内存版本接受以下 URI 格式：

filename?format=libsvm#cacheprefix

filename 是你想要加载的LIBSVM格式文件的常规路径，而 cacheprefix 是XGBoost将用于缓存预处理数据以二进制形式的缓存文件的路径。

要从csv文件加载，请使用以下语法：

filename.csv?format=csv&label_column=0#cacheprefix

其中 label_column 应指向作为标签的 csv 列。

如果你有一个以LIBSVM格式存储在类似 demo/data/agaricus.txt.train 文件中的数据集，可以通过以下方式启用外部内存支持：

dtrain = DMatrix('../data/agaricus.txt.train?format=libsvm#dtrain.cache')

XGBoost 将首先加载 agaricus.txt.train，对其进行预处理，然后将其写入名为 dtrain.cache 的新文件中，作为存储预处理数据在内部二进制格式中的磁盘缓存。有关文本输入格式的更多说明，请参阅 DMatrix 的文本输入格式。

对于CLI版本，只需添加缓存后缀，例如 "../data/agaricus.txt.train?format=libsvm#dtrain.cache"。