大型 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 在内存中分配了一个新对象。 （试图一次生成所有特征向量，然后连接会炸毁内存）。

也就是说，在循环结束附近运行串联仍然需要大约 30 GB 内存（在串联运行后立即释放）。

基本上，我的实例上有足够的内存来存储完整的功能集，但是内存从将东西打包到单个数组中跳跃使我内存不足。

有没有更节省内存的方法来运行它？

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)

您需要在np.fromiter 中指定dtype，并且可以选择使用count 指定从生成器获取的元素数量。虽然它是可选的，但最好指定 count 以便 numpy 可以预先分配输出数组，而不是根据需要调整它的大小。