大型Numpy数组的高效并列化

Question

我正在运行一个创建大量特征向量（作为numpy数组）并将它们堆叠成一个数组的过程。这个过程目前非常耗费内存，我正在寻找一种更有效的方式来运行它。

目前，我以100000个批次生成特征向量，并将它们连在一起。

all_features = None

for i in range(0, num_entries, 100000):
    features = get_features(entries[i:i+100000]) # generate a batch of 100,000 feature vectors
    features = np.array(features)

    if all_features is not None:
        all_features = np.concatenate([all_features, features])
    else:
        all_features = features

    del features
    gc.collect()

我发现，迭代连接特征向量，然后删除中间的 features 对象比一次生成所有的特征并将它们一次性连接起来更节省内存。我相信这是因为 np.concatenate 在内存中分配一个新对象。(试图一次生成所有的特征向量，然后再进行连接，会炸掉内存)。

也就是说，到了一个点，在循环接近尾声的时候，运行连接仍然需要大约30GB的内存（在连接运行后，这些内存立即被释放）。

基本上我的实例上有足够的内存来存储完整的功能集，但是把东西打包成一个数组后的内存跳跃让我的内存耗尽。

有没有一种更有效的内存运行方式？

Answer 1

如果知道all_features的总大小，我建议提前分配它。all_features=np.zeros(...) 然后在循环中填充它。这样，你就摆脱了多次重分配、删除和np.concatenate()调用。

Answer 2

让你的 get_features 函数生成器，然后用 np.fromiter 来创建数组。

简单的例子

def gen_values():
    for i in range(1000000): 
        yield i

a = np.fromiter(gen_values(), dtype=int)

您需要指定 dtype 在 np.fromiter 你还可以选择指定从生成器中获取的元素数量，用 count. 虽然它是可选的，但最好指定为 count 这样numpy就可以预先分配输出数组，而不是按需调整其大小。