Python在for循环中不释放内存

Question

我正在开发一个函数来对大量输入数据进行处理。由于我无法一次将所有数据装入内存（117703x200000矩阵用于点积），因此我将其分成块并分别计算。

输出仅采用排序后的前5个元素，因此必须为117703x5形状，这可以保持在内存中。但是，由于循环执行，我的内存消耗不断增加，直到出现内存错误。你有任何想法吗？以下是代码：

def process_predictions_proto(frac=50):
    # Simulate some inputs
    query_embeddings = np.random.random((117703, 512))
    proto_feat = np.random.random((200000, 512))
    gal_cls = np.arange(200000)

    N_val = query_embeddings.shape[0]
    pred = []

    for i in tqdm(range(frac)):
        start = i * int(np.ceil(N_val / frac))
        stop = (i + 1) * int(np.ceil(N_val / frac))
        val_i = query_embeddings[start:stop, :]
        # Compute distances
        dist_i = np.dot(val_i, proto_feat.transpose())
        # Sort
        index_i = np.argsort(dist_i, axis=1)[::-1]
        dist_i = np.take_along_axis(dist_i, index_i, axis=1)
        # Convert distances to class_ids
        pred_i = np.take_along_axis(
            np.repeat(gal_cls[np.newaxis, :], index_i.shape[0], axis=0),
            index_i, axis=1)
        # Use pd.unique to remove copies of the same class_id and
        # get 5 most similar ids
        pred_i = [pd.unique(pi)[:5] for pi in pred_i]
        # Append to list
        pred.append(pred_i)
        # Free memory
        gc.collect()
    pred = np.stack(pred, 0)  # N_val x 5
    return pred

- Eduardo

Answer 1

在调用gc.collect()之前删除所有临时变量，这样数据将立即成为垃圾。

del start, stop, val_i, dist_i, index_i, dist_i, pred_i
gc.collect()

在你的代码中，当你第一次调用gc.collect()时，没有任何数据被清理，因为所有变量都还能引用它们。第一轮迭代的数据直到第二轮结束后才会被清理；在第一轮之后的每一轮迭代中，内存中会有两个数据块（当前迭代和前一个迭代）。因此，你使用了比实际需要多一倍的内存（我假设某些对象之间存在引用，因此自动垃圾回收在循环期间重新分配变量时并不清理对象）。

- Barmar

---

那么，在循环开头简单地放置“gc.collect()”就足够了吗？ - gebbissimo

不行，因为那样的话，在最后一个迭代中分配的对象将不会被回收。 - Barmar