【发布时间】:2021-03-18 19:55:05
【问题描述】:
我正在做一个项目,我需要找到嵌入向量的最近邻。最近,我正在尝试使用 Google 的新 ANN 工具 SCANN github。 我能够创建搜索器对象并将其序列化为一个小数据集(约 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')
但是当我尝试使用真实数据集(约 48M 行,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 消息。我还是不知道为什么会这样……
[edit2] 我从源代码构建了 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
【问题讨论】:
-
一个明显的问题。您使用的是 32 位还是 64 位版本的 Python? 32 位版本只能使用 2GB 内存。
-
platform.architecture()的输出是('64bit', 'ELF')。我认为这可能是问题所在,但我刚刚成功序列化了一个 5M 行和 11GB 的数据集,但失败了 10M 行和 21GB -
在 Linux 上很难获得
bad_alloc。 Linux 通常会假装有足够的内存,即使没有。这称为过度使用,由proc/sys/vm/overcommit_memory控制。使用不可用内存的结果是 Linux 首先变慢(交换),然后杀死一个随机进程。bad_alloc更多地提示编程错误,例如尝试分配 -1 字节。 -
@MSalters 我已经进行了一些测试,似乎如果我的数据集的大小小于 16777126 (2^24) 就可以了。所以我猜你是对的,在 SCANN 代码的某个地方,要分配的内存大小溢出了一个变量,然后试图分配一个负数的内存。现在我只需要找出在哪里
-
看起来像 SCANN 或它的使用方式中的一个错误(不是这方面的专家,所以不能肯定地说)。
hash_dim和expected_size的类型都是int,所以hash_dim * expected_size溢出。size_t或int64_t会更好。
标签: python c++ python-3.x numpy out-of-memory