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

Answer 1

Удалите все временные переменные перед вызовом gc.collect(), чтобы данные сразу стали мусором.

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

В вашем коде, когда вы вызываете gc.collect() в первый раз, ни один из данных не является мусором, потому что на него все еще можно ссылаться из всех переменных. Данные из первой итерации не будут собраны до конца второй итерации; во время каждой итерации после первой вы будете иметь в памяти два фрагмента данных (текущую и предыдущую итерацию). Таким образом, вы используете вдвое больше памяти, чем вам нужно (я предполагаю, что между некоторыми объектами есть ссылки, поэтому автоматический сборщик мусора не очищает объекты, так как переменные переназначаются во время цикла).