NumPy 中的記憶體管理#

numpy.ndarray 是一個 python 類別。它需要額外的記憶體配置來保存 numpy.ndarray.stridesnumpy.ndarray.shapenumpy.ndarray.data 屬性。這些屬性是在 __new__ 中建立 python 物件後特別分配的。stridesshape 儲存在內部配置的一塊記憶體中。

用於儲存實際陣列值的 data 配置(在 object 陣列的情況下可能是指標)可能非常大,因此 NumPy 提供了介面來管理其配置和釋放。本文檔詳細說明了這些介面的運作方式。

歷史概觀#

自 1.7.0 版本以來,NumPy 公開了一組 PyDataMem_* 函數 (PyDataMem_NEWPyDataMem_FREEPyDataMem_RENEW),這些函數分別由 allocfreerealloc 支援。

自早期以來,Python 也改進了其記憶體管理功能,並從 3.4 版本開始提供各種 管理策略。這些常式分為一組域,每個域都有一個用於記憶體管理的 PyMemAllocatorEx 常式結構。Python 還新增了一個 tracemalloc 模組來追蹤對各種常式的呼叫。這些追蹤掛鉤已新增到 NumPy PyDataMem_* 常式中。

NumPy 在其內部 npy_alloc_cachenpy_alloc_cache_zeronpy_free_cache 函數中新增了一個小型已配置記憶體快取。這些函數分別封裝了 allocalloc-and-memset(0)free,但是當呼叫 npy_free_cache 時,它會將指標新增到按大小標記的可用區塊的簡短清單中。這些區塊可以由後續對 npy_alloc* 的呼叫重複使用,從而避免記憶體抖動。

NumPy 中的可配置記憶體常式 (NEP 49)#

使用者可能希望使用自己的常式覆寫內部資料記憶體常式。由於 NumPy 沒有使用 Python 域策略來管理資料記憶體,因此它提供了一組替代的 C-API 來變更記憶體常式。對於大量的物件資料,沒有 Python 域範圍的策略,因此這些策略不太適合 NumPy 的需求。希望變更 NumPy 資料記憶體管理常式的使用者可以使用 PyDataMem_SetHandler,它使用 PyDataMem_Handler 結構來保存用於管理資料記憶體的函數的指標。呼叫仍然由內部常式封裝,以呼叫 PyTraceMalloc_TrackPyTraceMalloc_Untrack。由於函數可能會在進程的生命週期內變更,因此每個 ndarray 都帶有在其實例化時使用的函數,這些函數將用於重新配置或釋放實例的資料記憶體。

type PyDataMem_Handler#

用於保存用於操作記憶體的函數指標的結構

typedef struct {
    char name[127];  /* multiple of 64 to keep the struct aligned */
    uint8_t version; /* currently 1 */
    PyDataMemAllocator allocator;
} PyDataMem_Handler;

其中 allocator 結構為

/* The declaration of free differs from PyMemAllocatorEx */
typedef struct {
    void *ctx;
    void* (*malloc) (void *ctx, size_t size);
    void* (*calloc) (void *ctx, size_t nelem, size_t elsize);
    void* (*realloc) (void *ctx, void *ptr, size_t new_size);
    void (*free) (void *ctx, void *ptr, size_t size);
} PyDataMemAllocator;
PyObject *PyDataMem_SetHandler(PyObject *handler)#

設定新的配置策略。如果輸入值為 NULL,將會將策略重設為預設值。傳回先前的策略,如果發生錯誤,則傳回 NULL。我們封裝了使用者提供的函數,以便它們仍然會呼叫 python 和 numpy 記憶體管理回呼掛鉤。

PyObject *PyDataMem_GetHandler()#

傳回將用於為下一個 PyArrayObject 配置資料的目前策略。失敗時,傳回 NULL

有關設定和使用 PyDataMem_Handler 的範例,請參閱 numpy/_core/tests/test_mem_policy.py 中的測試

如果未設定策略,則在解除配置時會發生什麼情況#

一種罕見但有用的技術是在 NumPy 外部配置緩衝區,使用 PyArray_NewFromDescr 將緩衝區包裝在 ndarray 中,然後將 OWNDATA 標誌切換為 true。當 ndarray 釋放時,應該從 ndarrayPyDataMem_Handler 中呼叫適當的函數以釋放緩衝區。但是 PyDataMem_Handler 欄位從未設定,它將為 NULL。為了向後相容性,NumPy 將呼叫 free() 以釋放緩衝區。如果 NUMPY_WARN_IF_NO_MEM_POLICY 設定為 1,則會發出警告。目前的預設值是不發出警告,這可能會在未來版本的 NumPy 中變更。

更好的技術是使用 PyCapsule 作為基礎物件

/* define a PyCapsule_Destructor, using the correct deallocator for buff */
void free_wrap(void *capsule){
    void * obj = PyCapsule_GetPointer(capsule, PyCapsule_GetName(capsule));
    free(obj);
};

/* then inside the function that creates arr from buff */
...
arr = PyArray_NewFromDescr(... buf, ...);
if (arr == NULL) {
    return NULL;
}
capsule = PyCapsule_New(buf, "my_wrapped_buffer",
                        (PyCapsule_Destructor)&free_wrap);
if (PyArray_SetBaseObject(arr, capsule) == -1) {
    Py_DECREF(arr);
    return NULL;
}
...

使用 np.lib.tracemalloc_domain 進行記憶體追蹤的範例#

請注意,自 Python 3.6(或更新版本)起,可以使用內建的 tracemalloc 模組來追蹤 NumPy 內部的配置。NumPy 將其 CPU 記憶體配置放入 np.lib.tracemalloc_domain 域中。如需更多資訊,請查看:https://docs.python.org/3/library/tracemalloc.html

以下是如何使用 np.lib.tracemalloc_domain 的範例

"""
   The goal of this example is to show how to trace memory
   from an application that has NumPy and non-NumPy sections.
   We only select the sections using NumPy related calls.
"""

import tracemalloc
import numpy as np

# Flag to determine if we select NumPy domain
use_np_domain = True

nx = 300
ny = 500

# Start to trace memory
tracemalloc.start()

# Section 1
# ---------

# NumPy related call
a = np.zeros((nx,ny))

# non-NumPy related call
b = [i**2 for i in range(nx*ny)]

snapshot1 = tracemalloc.take_snapshot()
# We filter the snapshot to only select NumPy related calls
np_domain = np.lib.tracemalloc_domain
dom_filter = tracemalloc.DomainFilter(inclusive=use_np_domain,
                                      domain=np_domain)
snapshot1 = snapshot1.filter_traces([dom_filter])
top_stats1 = snapshot1.statistics('traceback')

print("================ SNAPSHOT 1 =================")
for stat in top_stats1:
    print(f"{stat.count} memory blocks: {stat.size / 1024:.1f} KiB")
    print(stat.traceback.format()[-1])

# Clear traces of memory blocks allocated by Python
# before moving to the next section.
tracemalloc.clear_traces()

# Section 2
#----------

# We are only using NumPy
c = np.sum(a*a)

snapshot2 = tracemalloc.take_snapshot()
top_stats2 = snapshot2.statistics('traceback')

print()
print("================ SNAPSHOT 2 =================")
for stat in top_stats2:
    print(f"{stat.count} memory blocks: {stat.size / 1024:.1f} KiB")
    print(stat.traceback.format()[-1])

tracemalloc.stop()

print()
print("============================================")
print("\nTracing Status : ", tracemalloc.is_tracing())

try:
    print("\nTrying to Take Snapshot After Tracing is Stopped.")
    snap = tracemalloc.take_snapshot()
except Exception as e:
    print("Exception : ", e)