C++在Python中出现内存不足，但仍有足够的空间

Question

我正在一个项目上工作，需要找到一个嵌入向量的最近邻。最近我尝试使用Google的新ANN工具SCANNGitHub。我已经成功创建了搜索对象，并对一个小数据集（约200K行，512个值）进行序列化，代码如下：

import numpy as np
import scann
data = np.random.random((200k,512))
data = data / np.linalg.norm(data, axis=1)[:, np.newaxis]
searcher = scann.scann_ops_pybind.builder(data, 10, "dot_product").tree(
    num_leaves=2000, num_leaves_to_search=100, training_sample_size=250000).score_ah(
    2, anisotropic_quantization_threshold=0.2).reorder(100).build()
searcher.serialize('./scann')

但是当我使用真实数据集（大约4800万行，每行有512个值）时，我得到了以下结果：

In [11]: searcher.serialize('scann/')
---------------------------------------------------------------------------
MemoryError                               Traceback (most recent call last)
<ipython-input-11-71a5ef71c81f> in <module>
----> 1 searcher.serialize('scann/')

~/.local/lib/python3.6/site-packages/scann/scann_ops/py/scann_ops_pybind.py in serialize(self, artifacts_dir)
     70
     71   def serialize(self, artifacts_dir):
---> 72     self.searcher.serialize(artifacts_dir)
     73
     74

MemoryError: std::bad_alloc

数据集的.npy文件大小约为90GB，我的计算机上还有至少500GB的可用RAM和1TB的可用磁盘空间:

我正在运行Ubuntu 18.04.5 LTS和Python 3.6.9。Scann模块是通过Pip安装的。

有什么想法吗？会出现什么问题吗？

谢谢帮助

[编辑] 在@MSalters的评论后，我进行了一些测试，并发现如果要序列化的数据集超过16777220字节（2 ^ 24 + 4），它会出现bad_alloc消息而失败。我仍然不知道为什么会这样...

[编辑2] 我从源代码中构建了SCANN，并在其中加入了一些调试打印。错误似乎在这一行：

vector<uint8_t> storage(hash_dim * expected_size);

如果我这样打印它：

std::cout << hash_dim <<  " " << expected_size <<"\n" << std::flush;
std::cout << hash_dim * expected_size <<"\n" << std::flush;
vector<uint8_t> v2;
std::cout << v2.max_size() << "\n" << std::flush;
vector<uint8_t> storage(hash_dim * expected_size);
std::cout << "after storage creation\n" << std::flush;

然后我获得；

256 8388608
-2147483648
9223372036854775807

- Feulo

---

有一个显而易见的问题。您正在使用32位还是64位版本的Python？32位版本只能使用2GB的内存。 - Simon Callan

platform.architecture() 的输出是 ('64bit', 'ELF')。我以为这可能是问题所在，但我刚刚成功地序列化了一个包含 5M 行和 11GB 的数据集，但在尝试序列化包含 10M 行和 21GB 的数据集时失败了。 - Feulo

在Linux上很难出现bad_alloc错误。即使内存不足，Linux通常会假装有足够的内存，这被称为过度承诺，并由proc/sys/vm/overcommit_memory控制。使用不可用内存的结果是，Linux首先变得缓慢（交换），然后杀死一个随机进程。bad_alloc更可能是编程错误的提示，例如尝试分配-1字节的空间。 - MSalters

@MSalters 我已经进行了一些测试，似乎如果我的数据集大小小于16777126（2 ^ 24），它就可以正常工作。所以我想你是对的，在SCANN代码的某个地方，要分配的内存大小超出了变量的范围，然后尝试分配负数内存量。现在我只需要找出在哪里。 - Feulo

看起来是 SCANN 中的 bug 或者使用方式有问题（不是专家，不能确定）。hash_dim 和 expected_size 都是 int 类型，所以 hash_dim * expected_size 溢出了。size_t 或 int64_t 更适合使用。 - rustyx

Answer 1

在ScaNN中似乎存在一个现有的问题报告#427，其中包含类似的错误。

根据std::cout << hash_dim * expected_size 的输出结果为-2147483648，我们可以得出结论，hash_dim * expected_size 溢出了。

查看源代码，我们发现 hash_dim 和 expected_size 的类型都是 int。

所以，这些中至少有一个应该是int64_t、long long或者更好的size_t类型。通过查看ScaNN的源代码，似乎还有更多地方可以受益于专门设计用于保存大小（size_t）而不是int的数据类型。

- rustyx