用 NumPy 中的視圖來節省內存
本文轉載自微信公眾號「Python中文社區」,作者Trauring。轉載本文請聯系Python中文社區公眾號。
如果您使用 Python 的 NumPy 庫,通常是因為您正在處理占用大量內存的大型數組。為了減少內存使用,您可能希望盡量減少不必要的重復項。
NumPy 有一個內置功能,可以在許多常見情況下透明地執行此操作:內存視圖。而且,此功能還可以防止數組被垃圾回收,從而導致更高的內存使用率。在某些情況下,它可能會導致錯誤,數據會以意想不到的方式發生變異。
為了避免這些問題,讓我們了解視圖的工作原理以及對代碼的影響。
預備知識:Python 列表
在查看 NumPy 數組和視圖之前,讓我們考慮一個有點相似的數據結構:Python 列表。
Python 列表與 NumPy 數組一樣,是連續的內存塊。當你對一個 Python 列表進行切片時,你會得到一個完全不同的列表,這意味著你正在分配一塊新的內存:
- >>> from psutil import Process
- >>> Process().memory_info().rss
- 12247040
- >>> list1 = [None] * 1_000_000
- >>> Process().memory_info().rss
- 20463616
- >>> list2 = list1[:500_000]
- >>> Process().memory_info().rss
- 24580096
切片列表分配了更多內存。由于第二個列表是一個獨立的副本,如果我們改變它,這不會影響第一個列表:
- >>> list2[0] = "abc"
- >>> print(list2[0])
- abc
- >>> print(list1[0])
- None
注意,復制到第二個列表中的數據是指向 Python 對象的指針,而不是對象本身的內容。因此,即使列表本身不同,底層對象仍然在兩者之間進行共享。
切片時 NumPy 數組并不進行復制
NumPy 數組的工作方式不同。因為假設您可能正在處理非常大的數組,所以許多操作不會復制數組,它們只是讓您查看原始數組指向的同一連續內存塊。
第一個結果是切片不會分配更多內存,因為它只是原始數組的視圖:
- >>> from psutil import Process
- >>> import numpy as np
- >>> arr = np.arange(0, 1_000_000)
- >>> Process().memory_info().rss
- 37810176
- >>> view = arr[:500_000]
- >>> Process().memory_info().rss
- 37810176
視圖對象看起來像一個 500,000 長的 int64 數組,因此如果它是一個新數組,它將分配大約 4MB 的內存。但它只是針對同一個原始數組的一個視圖,所以不需要額外的內存。
從技術上來說,可能會為視圖對象本身分配一小部分內存,但這可以忽略不計,除非您有很多視圖對象。在這種情況下,RSS(常駐內存)度量中沒有出現新內存,因為 Python 預先分配了更大的內存塊,然后用小的 Python 對象填充這些塊。
視圖導致內存泄漏
使用視圖的后果之一是您可能會泄漏內存,而不是節省內存。這是因為視圖可以防止原始數組被垃圾回收 - 對整個數組來說。
假設您已經決定只需要使用大數組的一小部分:
- >>> import numpy as np
- >>> from psutil import Process
- >>> arr = np.arange(0, 100_000_000)
- >>> Process().memory_info().rss
- 830181376
- >>> small_slice = arr[:10]
- >>> del arr
- >>> Process().memory_info().rss
- 830111744
如果這是一個 Python 列表,刪除原始對象將釋放內存。然而,在這種情況下,即使我們沒有對數組的直接引用,視圖仍然可以起作用,這意味著內存沒有被釋放,即使我們只對其中的一小部分感興趣。
您實際上可以通過視圖訪問原始數組:
- >>> small_slice
- array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
- >>> small_slice.base
- array([0, 1, 2, ..., 99999997, 99999998, 99999999])
結果,只有當我們刪除所有視圖時,原始數組的內存才會被釋放:
- >>> del small_slice
- >>> Process().memory_info().rss
- 29642752
其他改變
使用視圖的另一個后果是修改視圖會改變原始數組。回想一下,對于 Python 列表,修改切片結果不會修改原始列表,因為新對象是一個副本:
- >>> l = [1, 2, 3]
- >>> l2 = l[:]
- >>> l2[0] = 17
- >>> l2
- [17, 2, 3]
- >>> l
- [1, 2, 3]
使用 NumPy 視圖后,改變視圖確實改變了原始對象,它們都指向同一個內存地址:
- >>> arr = np.array([1, 2, 3])
- >>> view = arr[:]
- >>> view[0] = 17
- >>> view
- array([17, 2, 3])
- >>> arr
- array([17, 2, 3])
這個結果不是我們想要的!
由于某些 NumPy API 可能會根據情況返回視圖或副本,因此更有可能發生意外變化。例如,某些切片結果可能不是視圖:
- >>> arr = np.array([1, 2, 3])
- >>> arr2 = arr[:]
- >>> arr2.base is arr
- True
- >>> arr3 = arr[[True, False, True]]
- >>> arr3.base is arr
- False
改變 arr2 也會改變 arr,但改變 arr3 不會改變 arr。
使用 copy() 進行顯式復制
當您不想引用原始內存時,顯式復制允許您創建一個新數組。這對于防止改變很有用,并且在您不想將原始數組保留在內存中的情況下也很有用:
- >>> arr = np.arange(0, 100_000_000)
- >>> Process().memory_info().rss
- 829464576
- >>> small_slice = arr[:10].copy()
- >>> del arr
- >>> Process().memory_info().rss
- 29700096
- >>> print(small_slice.base)
- None
在這種情況下,刪除 arr 釋放了內存,因為 small_slice 是副本,而不是視圖。
要點:高效安全地使用視圖
鑒于各種 NumPy API 會自動返回視圖,您需要在編寫代碼時考慮它們:
- 在文檔中注意 API 是否會返回視圖、副本或兩者。
- 如果您想從內存中清除一個大數組,請確保不僅沒有直接引用它,而且沒有引用它的視圖。
- 如果你要改變一個數組,確保它不會因為它實際上是一個視圖而意外改變其他一些數組。
- 如果您不需要視圖,請使用 copy() 方法。