有哪些方法可以加快 Pytorch 中大型稀疏数组（约 100 万 x 100 万，密度约 0.0001）的数据加载速度？

Posted 2023-04-18

【中文标题】有哪些方法可以加快 Pytorch 中大型稀疏数组（约 100 万 x 100 万，密度约 0.0001）的数据加载速度？

【英文标题】：What are some ways to speed up data loading on large sparse arrays (~1 million x 1 million, density ~0.0001) in Pytorch?

【发布时间】：2020-07-22 15:14:47

【问题描述】：

我正在研究一个二元分类问题。我有大约 150 万个数据点，特征空间的维数是 100 万。该数据集存储为稀疏数组，密度约为 0.0001。对于这篇文章，我将限制范围假设模型是浅层前馈神经网络，并假设维度已经优化（因此不能减少到 100 万以下）。从这些数据中创建小批量以馈送到网络的简单方法将花费大量时间（例如，从输入数组的torch.sparse.FloatTensor 表示创建TensorDataset（地图样式）的基本方法，并在其周围包裹一个DataLoader，这意味着~20s 可以获得一个 32 的 mini-batch 到网络，而不是说 ~0.1s 来执行实际训练）。我正在寻找加快速度的方法。

我尝试过的

我首先想到，在 DataLoader 的每次迭代中从如此大的稀疏数组中读取是计算密集型的，因此我将这个稀疏数组分解为更小的稀疏数组 为了让DataLoader 以迭代方式从这些多个稀疏数组中读取数据，我将DataLoader 中的地图样式数据集替换为IterableDataset，并将这些较小的稀疏数组流式传输到这个IterableDataset 中，例如所以：

from itertools import chain
from scipy import sparse

class SparseIterDataset(torch.utils.data.IterableDataset):
    
    def __init__(self, fpaths):
        super(SparseIter).__init__()
        self.fpaths = fpaths
    
    def read_from_file(self, fpath):
        data = sparse.load_npz(fpath).toarray()
        for d in data:
            yield torch.Tensor(d)
            
    def get_stream(self, fpaths):
        return chain.from_iterable(map(self.read_from_file, fpaths))
    
    def __iter__(self):
        return self.get_stream(self.fpaths)

通过这种方法，我能够将时间从大约 20 秒的幼稚基本情况缩短到每 32 个小批量的大约 0.2 秒。但是，鉴于我的数据集有大约 150 万个样本，这仍然意味着很多甚至一次通过数据集所花费的时间。 （作为比较，即使它有点苹果对橘子，在原始稀疏数组上对 scikit-learn 运行逻辑回归大约需要 6 秒，每次迭代通过整个数据集。使用 pytorch，使用方法我刚刚概述了，仅在一个 epoch 中加载所有 minibatch 大约需要 3000 秒）

我知道但尚未尝试的一件事是通过在DataLoader 中设置num_workers 参数来使用多进程数据加载。我相信这在可迭代样式数据集的情况下有它自己的问题。另外，即使是 10 倍的加速，仍然意味着每个 epoch 加载小批量约 300 秒。我觉得我的速度太慢了！您是否可以提出任何其他方法/改进/最佳实践建议？

【问题讨论】：

坦率地说，这不是很多数据——它是什么，150m 的值？最大的瓶颈将是toarray()，它必须分配大量的内存来保存所有这些零。为什么不直接将 coo 索引推送到 torch.sparse_coo_tensor 中？我认为没有任何理由必须对这个数组进行加密。

【参考方案1】：

未稀疏形式的数据集将是 1.5M x 1M x 1 字节 = 1.5TB 作为 uint8，或 1.5M x 1M x 4 字节 = 6TB 作为 float32。在现代 CPU 上，简单地从内存读取 6TB 到 CPU 可能需要 5-10 分钟（取决于架构），并且从 CPU 到 GPU 的传输速度会比这慢一些（PCIe 上的 NVIDIA V100 理论上有 32GB/s）。

方法：

单独对所有内容进行基准测试 - 例如在 jupyter 中

%%timeit data = sparse.load_npz(fpath).toarray()

%%timeit dense = data.toarray() # 去稀疏化比较

%%timeit t = torch.tensor(data) # 大概和上面那行差不多

还打印出所有内容的形状和数据类型，以确保它们符合预期。我还没有尝试运行您的代码，但我很确定（a）sparse.load_npz 非常快并且不太可能成为瓶颈，但是（b）torch.tensor(data) 产生密集张量并且在这里也很慢

使用torch.sparse。我认为在大多数情况下，torch 稀疏张量可以用作常规张量。您必须做一些数据准备才能从 scipy.sparse 转换为 torch.sparse：

稀疏张量表示为一对密集张量： 值和索引的二维张量。 稀疏张量可以由构造 提供这两个张量，以及稀疏张量的大小

您提到了torch.sparse.FloatTensor，但我很确定您没有在代码中创建稀疏张量 - 没有理由期望这些会简单地通过将 scipy.sparse 数组传递给常规张量构造函数来构造，因为它们通常不是这样制作的。

如果你想出一个好的方法，我建议你将它作为项目或 git 发布在 github 上，这将非常有用。

如果 torch.sparse 不起作用，请考虑其他方法，将数据仅在 GPU 上转换为密集数据，或避免完全转换。

另请参阅： https://towardsdatascience.com/sparse-matrices-in-pytorch-be8ecaccae6 https://github.com/rusty1s/pytorch_sparse

【讨论】：

感谢您抽出宝贵时间回复！对于第 2 点，我没有将 scipy.sparse 数组传递给张量构造函数，而是将其值和索引传递给（这是推荐的方式）——为简洁起见，我省略了这一点。您发布的链接看起来很有趣。在接下来的几天里，我将继续探讨这个问题。

@AbhimanyuSahai - 如何构造torch.sparse.FloatTensor 是最重要的部分，所以请添加该代码。您显示的代码示例将稀疏转换为密集，这听起来不是您想要的。仅此一项就可以解释速度慢的原因。