使用“不安全的 Python”加速 Numpy 代碼 100 倍！
共 14410字，需瀏覽 29分鐘
·
2024-05-28 22:21
△△請給“Python貓”加星標 ，以免錯過文章推送
原文：A 100x speedup with unsafe Python[1]
我們將使用“不安全”的Python將一些Numpy代碼加速100倍。假設(shè)你在用pygame編寫一個游戲，并且你需要經(jīng)常調(diào)整圖像大小。我們可以使用pygame或openCV調(diào)整圖像大小：
from contextlib import contextmanager
import time
import pygame as pg
import numpy as np
import cv2
@contextmanager
def Timer(name):
start = time.time()
yield
finish = time.time()
print(f'{name} took {finish-start:.4f} sec')
IW = 1920
IH = 1080
OW = IW // 2
OH = IH // 2
repeat = 10 # 譯者注：原文此處為100，參考作者github代碼和后續(xù)運行結(jié)果，應(yīng)該是10。
isurf = pg.Surface((IW,IH), pg.SRCALPHA)
with Timer('pg.Surface with smoothscale'):
for i in range(repeat):
pg.transform.smoothscale(isurf, (OW,OH))
def cv2_resize(image):
return cv2.resize(image, (OH,OW), interpolation=cv2.INTER_AREA)
i1 = np.zeros((IW,IH,3), np.uint8)
with Timer('np.zeros with cv2'):
for i in range(repeat):
o1 = cv2_resize(i1)輸出為：
pg.Surface with smoothscale took 0.2002 sec
np.zeros with cv2 took 0.0145 sec使用openCV加速了10倍，所以我們回到游戲中，使用pygame.surfarray.pixels3d以零拷貝的方式訪問像素作為Numpy數(shù)組，然后使用cv2.resize，然而，一切都變慢了！
i2 = pg.surfarray.pixels3d(isurf)
with Timer('pixels3d with cv2'):
for i in range(repeat):
o2 = cv2_resize(i2)結(jié)果：
pixels3d with cv2 took 1.3625 sec如果你去查看數(shù)組的.shape和.dtype，會發(fā)現(xiàn)他倆一樣。但是，同一個函數(shù)(cv2_resize)在一個數(shù)組上運行比另一個數(shù)組慢 100 倍，為什么捏？SDL 應(yīng)該不會在一些特別難以訪問的 RAM 區(qū)域分配像素（即使它在理論上可以通過內(nèi)核的一點幫助來做到這一點，比如創(chuàng)建一個不可緩存的內(nèi)存區(qū)域之類的）?；蛘逽urface存儲在 GPU 中，通過 PCI 獲取每個像素？！它不是這樣工作的，是嗎？-這些東西有一些可怕的內(nèi)存一致性協(xié)議，我錯過了什么嗎？如果不是——如果它們是相同形狀和大小的相同類型的內(nèi)存——是什么不同導(dǎo)致我們減速 100 倍？
結(jié)果證明...我承認我是偶然發(fā)現(xiàn)的，在放棄這個并轉(zhuǎn)向其他事情之后。完全偶然的是，那個其他事情涉及將 numpy 數(shù)據(jù)傳遞給 C 代碼，所以我不得不學(xué)習(xí)這個數(shù)據(jù)在 C 中的樣子。所以，事實證明，shape和datatype并不是 numpy 數(shù)組的全部：
print('input strides',i1.strides,i2.strides)
print('output strides',o1.strides,o2.strides)結(jié)果，輸出的步幅(stride)是不同的：
input strides (3240, 3, 1) (4, 7680, -1)
output strides (1620, 3, 1) (1620, 3, 1)所以是步幅的差異導(dǎo)致了代碼慢了100倍？我們可以修復(fù)這個問題嗎？但首先我們要去解釋為什么步幅不同。
numpy array內(nèi)存布局
所以步幅(stride)是什么？步幅告訴您從一個像素到下一個像素需要跨越多少字節(jié)。例如，假設(shè)我們有一個三維數(shù)組，比如一個 RGB 圖像。然后，給定數(shù)組的基指針和三個步幅， array[x,y,z]的地址將是 （對于圖像，z 的值為 0、1 或 2，分別對應(yīng) RGB 圖像的三個通道之一）。
換句話說，步幅定義了數(shù)組在內(nèi)存中的布局。無論好壞，numpy 在數(shù)組形狀和數(shù)據(jù)類型方面非常靈活，因為它支持許多不同的步幅值。
手頭的兩種布局 : numpy 的默認布局和 SDL 的布局 - 嗯，我甚至不知道哪個更冒犯我。從步幅值可以看出，numpy 默認用于 3D 數(shù)組的布局是 。
這意味著一個像素的 RGB 值存儲在 3 個相鄰的字節(jié)中，一列的像素在內(nèi)存中連續(xù)存儲 - 以列為主序。我覺得這種方法很冒犯，因為圖像傳統(tǒng)上是以行為主序存儲的，尤其是圖像傳感器以這種方式發(fā)送圖像（并以這種方式捕捉圖像，正如您可以從滾動快門看到的 - 每一行在稍微不同的時間點進行捕捉，而不是按列進行）
“為什么，我們確實也遵循這一受人尊敬的傳統(tǒng)，”流行的基于numpy的圖像庫說道?！翱纯茨阕约骸獙⒁粋€形狀為 (1920, 1080) 的數(shù)組保存為 PNG 文件，你會得到一張 1080x1920 的圖像”。這是真的，當然這使情況變得更糟：如果你使用 arr[x,y] 進行索引，那么 x，也就是維度零，實際上對應(yīng)于相應(yīng) PNG 文件中的垂直維度；而 y，也就是維度一，對應(yīng)于水平維度。因此，numpy 數(shù)組的列對應(yīng)于 PNG 圖像的行。這在某種意義上使 numpy 圖像布局成為"行優(yōu)先"，但代價是 x 和 y 的含義與通常相反。
...除非你從 pygame Surface 對象中獲取 numpy 數(shù)組，否則 x 實際上是索引到水平維度的。因此，相對于 pygame.image.save(surface) 創(chuàng)建的 PNG 文件，使用 imageio 保存 pixels3d(surface) 將會產(chǎn)生一個轉(zhuǎn)置的 PNG。而且，如果這種侮辱還不夠，cv2.resize 使用 (width, height) 元組作為目標大小，將產(chǎn)生一個形狀為 (height, width) 的輸出數(shù)組。
在這些侮辱和傷害的背景下，SDL 擁有一個誘人的、文明的布局，其中 x 是 x，y 是 y，數(shù)據(jù)以誠實的行優(yōu)先順序存儲，對于“行”的所有含義都是如此。但是仔細一看，這個布局只是踐踏了我的感情： 。
像是我們在步幅中有 4 而不是 3 的部分，對于 RGB 圖像我可以理解。當我們將 SRCALPHA 傳遞給 Surface 構(gòu)造函數(shù)時，我們確實要求一個帶有 alpha 通道的 RGBA 圖像。所以我猜它將 alpha 通道與 RGB 像素一起保留，并且步幅中的 4 需要跳過 RBGA 中的 A。但是，我想問，為什么有單獨的 pixels3d 和 pixels_alpha 函數(shù)？在使用 numpy 和 pygame Surface時，分別處理 RGB 和 alpha 總是很麻煩。為什么不是一個單一的 pixels4d 函數(shù)呢？
...好吧，4 而不是 3 我可以接受。但是 zstride 為-1？負一？你從紅色像素的地址開始，要到綠色，你要往回走一個字節(jié)？！現(xiàn)在你只是在拿我開玩笑。
原來 SDL 支持 RGB 和 BGR 布局（特別是，顯然從文件加載的surface是 RGB，而在內(nèi)存中創(chuàng)建的surface是 BGR？..或者比這更復(fù)雜？..）如果您使用 pygame 的 API，則無需擔心 RGB 與 BGR，API 會透明地處理它。如果您使用 pixels3d 進行 numpy 互操作，您也無需擔心 RGB 與 BGR，因為 numpy 的步幅靈活性使 pygame 可以為您提供一個看起來像 RGB 的數(shù)組，盡管它在內(nèi)存中是 BGR。為此，z 步幅設(shè)置為-1，并且數(shù)組的基指針指向第一個像素的紅色值-比數(shù)組內(nèi)存開始的位置提前兩個像素，即第一個像素的藍色值所在的位置。
等一下......現(xiàn)在我明白為什么我們有 pixels3d 和 pixels_alpha，但沒有 pixels4d 了??！因為 SDL 有 RGBA 和 BGRA 圖像——BGRA，而不是 ABGR——你無法使 BGRA 數(shù)據(jù)看起來像一個 RGBA numpy 數(shù)組，無論你使用怎樣奇怪的步幅值。布局靈活性是有限的......或者更確切地說，實際上沒有任何限制超過可計算限制，但幸運的是 numpy 止步于可配置步幅，并不允許您為完全可編程的布局 指定一個通用回調(diào)函數(shù) addr(base, x, y, z) 。
為了透明地支持 RGBA 和 BGRA，pygame 被迫給我們提供 2 個 numpy 數(shù)組 - 一個用于 RGB（或 BGR，取決于surface），另一個用于 alpha 通道。這些 numpy 數(shù)組具有正確的形狀，并讓我們訪問正確的數(shù)據(jù)，但它們的布局與其形狀的普通數(shù)組非常不同。
不同的內(nèi)存布局肯定可以解釋性能上的主要差異。我們可以試圖弄清楚為什么性能差異幾乎是 100 倍。但是如果可能的話，我更愿意擺脫垃圾，而不是詳細研究它。所以，我們不會深入理解這個問題，而是簡單地展示布局差異確實解釋了 100 倍的差異，然后在不改變布局的情況下擺脫減速，這就是“不安全的 Python”最終發(fā)揮作用的地方。
如何證明僅僅是布局，而不是 pygame Surface 數(shù)據(jù)的其他屬性（比如分配的內(nèi)存）導(dǎo)致了減速？我們可以對一個我們自己創(chuàng)建的具有與 pixels3d 相同布局的 numpy 數(shù)組進行 cv2.resize 的基準測試。
# create a byte array of zeros, and attach
# a numpy array with pygame-like strides
# to this byte array using the buffer argument.
i3 = np.ndarray((IW, IH, 3), np.uint8,
strides=(4, IW*4, -1),
buffer=b'\0'*(IW*IH*4),
offset=2) # start at the "R" of BGR
with Timer('pixels3d-like layout with cv2'):
for i in range(repeat):
o2 = cv2_resize(i3)確實，這幾乎和我們在 pygame Surface 數(shù)據(jù)上測量的一樣慢：
pixels3d-like layout with cv2 took 1.3829 sec出于好奇，我們可以檢查如果僅僅在這些布局之間復(fù)制數(shù)據(jù)會發(fā)生什么：
i4 = np.empty(i2.shape, i2.dtype)
with Timer('pixels3d-like copied to same-shape array'):
for i in range(repeat):
i4[:] = i2
with Timer('pixels3d-like to same-shape array, copyto'):
for i in range(repeat):
np.copyto(i4, i2)賦值運算符和 copyto 都非常慢，幾乎和 cv2.resize 一樣慢。
pixels3d-like copied to same-shape array took 1.2681 sec
pixels3d-like to same-shape array, copyto took 1.2702 sec愚弄代碼以使其運行更快
我們能做什么？我們無法改變 pygame Surface 數(shù)據(jù)的布局。我們也絕對不想復(fù)制 cv2.resize 的 C++代碼，因為它具有各種平臺特定的優(yōu)化，看看我們是否能夠適應(yīng) Surface 布局而不會丟失性能。我們可以做的是使用帶有 numpy 默認布局的數(shù)組將 Surface 數(shù)據(jù)饋送給 cv2.resize（而不是直接傳遞由 pixel3d 返回的數(shù)組對象）。
請注意，這實際上并不適用于任何給定的函數(shù)。但它將特別適用于調(diào)整大小，因為它實際上并不關(guān)心數(shù)據(jù)的某些方面，我們實際上會公然歪曲：
? 調(diào)整大小的代碼不在乎特定通道代表紅色還是藍色。（與將 RGB 轉(zhuǎn)換為灰度不同，后者會在意。）如果您給出 BGR 數(shù)據(jù)并謊稱它是 RGB，則代碼將產(chǎn)生與給出實際 RGB 數(shù)據(jù)時相同的結(jié)果。
? 同樣，調(diào)整大小時，數(shù)組維度代表寬度和高度的順序并不重要。
現(xiàn)在，讓我們再次來看看 pygame 的 BGRA 數(shù)組的內(nèi)存表示，其形狀是 (width, height) 。
這個表示實際上與一個形狀為 (height, width) 的 RGBA 數(shù)組具有 numpy 的默認步幅是一樣的！我的意思是，不完全一樣 - 如果我們將這個數(shù)據(jù)重新解釋為 RGBA 數(shù)組，我們將紅色通道(R)的值視為藍色(B)，反之亦然。同樣地，如果我們將這個數(shù)據(jù)重新解釋為一個具有 numpy 的默認步幅的 (height, width) 數(shù)組，我們將隱式地對圖像進行轉(zhuǎn)置。但是調(diào)整大小并不在乎！
而且，作為額外的好處，我們將得到一個單獨的 RGBA 數(shù)組，并且只需要一次調(diào)用 cv2.resize 來調(diào)整大小，而不是分別調(diào)整 pixels3d 和 pixels_alpha。耶！
下面的一段代碼接收一個 Pygame surface并返回底層 RGBA 或 BGRA 數(shù)組的基礎(chǔ)指針，以及一個指示它是 BGR 還是 RGB 的標志
import ctypes
def arr_params(surface):
pixels = pg.surfarray.pixels3d(surface)
width, height, depth = pixels.shape
assert depth == 3
xstride, ystride, zstride = pixels.strides
oft = 0
bgr = 0
if zstride == -1: # BGR image - walk back
# 2 bytes to get to the first blue pixel
oft = -2
zstride = 1
bgr = 1
# validate our assumptions about the data layout
assert xstride == 4
assert zstride == 1
assert ystride == width*4
base = pixels.ctypes.data_as(ctypes.c_void_p)
ptr = ctypes.c_void_p(base.value + oft)
return ptr, width, height, bgr既然我們獲得了像素數(shù)據(jù)的基礎(chǔ) C 指針，我們可以使用默認步長將其包裝在一個 numpy 數(shù)組中，隱式轉(zhuǎn)置圖像并交換 R＆B 通道。一旦我們將帶有默認步長的 numpy 數(shù)組“附加”到輸入和輸出數(shù)據(jù)上，我們對 cv2.resize 的調(diào)用將快 100 倍！
def rgba_buffer(p, w, h):
# attach a WxHx4 buffer to the base pointer
type = ctypes.c_uint8 * (w * h * 4)
return ctypes.cast(p, ctypes.POINTER(type)).contents
def cv2_resize_surface(src, dst):
iptr, iw, ih, ibgr = arr_params(src)
optr, ow, oh, obgr = arr_params(dst)
# our trick only works if both surfaces are BGR,
# or they're both RGB. if their layout doesn't match,
# our code would actually swap R & B channels
assert ibgr == obgr
ibuf = rgba_buffer(iptr, iw, ih)
# numpy's default strides are height*4, 4, 1
iarr = np.ndarray((ih,iw,4), np.uint8, buffer=ibuf)
obuf = rgba_buffer(optr, ow, oh)
oarr = np.ndarray((oh,ow,4), np.uint8, buffer=obuf)
cv2.resize(iarr, (ow,oh), oarr, interpolation=cv2.INTER_AREA)果然，我們最終從使用 cv2.resize 對 Surface 數(shù)據(jù)進行調(diào)整中獲得了加速而不是減速，我們的速度與調(diào)整 RGBA numpy.zeros 數(shù)組相同（最初我們對 RGB 數(shù)組進行基準測試，而不是 RGBA）
osurf = pg.Surface((OW,OH), pg.SRCALPHA)
with Timer('attached RGBA with cv2'):
for i in range(repeat):
cv2_resize_surface(isurf, osurf)
i6 = np.zeros((IW,IH,4), np.uint8)
with Timer('np.zeros RGBA with cv2'):
for i in range(repeat):
o6 = cv2_resize(i6) 基準測試顯示我們獲得了 100 倍的回報：
attached RGBA with cv2 took 0.0097 sec
np.zeros RGBA with cv2 took 0.0066 sec上面所有丑陋的代碼都在 GitHub[2] 上。由于這些代碼很丑陋，你不能確定它是否正確地調(diào)整了圖像大小，因此還有一些代碼在那里測試非零圖像的調(diào)整大小。如果你運行它，你將得到以下華麗的輸出圖像：
我們真的獲得了 100 倍的加速嗎？這取決于你如何計算。相對于直接使用 pixel3d 數(shù)組調(diào)用它，我們使 cv2.resize 的運行速度提高了 100 倍。但是特別是對于調(diào)整大小，pygame 有 smoothscale，相對于它，我們的加速比是 13-15 倍。在 GitHub 上還有一些其他函數(shù)的基準測試，其中一些沒有相應(yīng)的 pygame API。
? copying with dis[:] = src : 28x
? Inverting with dst[:] 255-src: 24x
? cv2.warpAffine: 12x
? cv2.mediaBlur: 15x
? cv2.GaussianBlur: 200x
無論如何，我會感到驚訝，如果有很多人使用 Python 從 SDL 來處理這個特定問題，以便使這個問題得到廣泛的關(guān)注。但我猜測，具有奇怪布局的 numpy 數(shù)組也可能在其他地方出現(xiàn)，因此這種技巧可能在其他地方也是相關(guān)的。
Unsafe Python
上面的代碼使用“C 風(fēng)格的知識”來加快速度（Python 通常會隱藏數(shù)據(jù)布局，而 C 則會自豪地暴露它。）不幸的是，它具有 C 的內(nèi)存（不）安全性 - 我們獲得了像素數(shù)據(jù)的 C 基指針，從那一點開始，如果我們搞砸了指針算術(shù)，或者在數(shù)據(jù)被釋放后繼續(xù)使用數(shù)據(jù)，我們就會崩潰或損壞數(shù)據(jù)。然而我們沒有編寫任何 C 代碼 - 這全部都是 Python。
Rust 有一個"unsafe"關(guān)鍵字，編譯器強制你意識到你正在調(diào)用一個會破壞正常安全性保證的 API。但是 Rust 編譯器并不會讓你把包含 unsafe 代碼塊的函數(shù)標記為"unsafe"。相反，它相信你可以決定你的函數(shù)本身是否安全。
在我們的示例中， cv2_resize_surface 是一個安全的 API，假設(shè)我沒有 Bug，因為沒有恐怖逃逸到外部世界 - 在外部，我們只看到輸出表面被輸出數(shù)據(jù)填充。但 arr_params 是一個完全不安全的 API，因為它返回一個 C 指針，你可以對其做任何操作。 rgba_buffer 也是不安全的——盡管我們返回一個 numpy 數(shù)組，一個“安全”的對象，但在數(shù)據(jù)被釋放后，你仍然可以使用它，例如。在一般情況下，沒有靜態(tài)分析可以告訴你是否從不安全的構(gòu)建模塊構(gòu)建了安全的東西。
Python 沒有 unsafe 關(guān)鍵字 - 這在動態(tài)語言和稀疏靜態(tài)注釋方面是符合特點的。但除此之外，Python + ctypes + C 庫在精神上有點類似于帶有 unsafe 的 Rust。該語言默認是安全的，但在需要時可以使用逃生通道。
《不安全的 Python》是一個通用原則的例證：Python 中大量使用了 C 語言。C 語言是 Python 的邪惡孿生兄弟，或者按時間順序來說，Python 是 C 語言的友好孿生兄弟。C 語言提供了性能，不關(guān)心可用性或安全性；如果其中任何一個導(dǎo)致問題，告訴你的醫(yī)療保健提供者，C 語言不感興趣。另一方面，Python 提供了安全性，并且基于十年來對初學(xué)者可用性的研究。不過，Python 不關(guān)心性能。它們都針對兩種相反的目標進行了激烈的優(yōu)化，忽視了對方的目標代價。
但更重要的是，Python 從一開始就考慮到了與 C 擴展的兼容性。今天，從我的角度來看，Python 作為一個流行的 C/C++ 庫的打包系統(tǒng)。我很少有下載和構(gòu)建 OpenCV 以在 C++ 中使用它的興趣，相較于使用 Python 中的 OpenCV 二進制文件，因為 C++ 沒有標準的包管理系統(tǒng)，而 Python 有。在 Python 中調(diào)用這些高性能庫（例如在科學(xué)計算和深度學(xué)習(xí)中）的代碼比在 C/C++ 中更多。另一方面，如果想要嚴格優(yōu)化的 Python 代碼和較小的部署文件大小/低啟動時間，你可以使用 Cython 來生成一個“仿寫成 C 所寫”的擴展，以節(jié)省類似 numba 這樣“更 Pythonic”的基于 JIT 的系統(tǒng)的開銷。
Python 中不僅有很多 C 代碼，而且它們是某種意義上的對立物，它們相互補充得相當好。使 Python 代碼快速的好方法是以正確的方式使用 C 庫。相反，安全使用 C 的好方法是用 C 編寫核心，然后在 Python 中編寫大量邏輯。Python 和 C/C++/Rust 混合——無論是具有大量 Python 擴展 API 的 C 程序，還是在 C 中完成所有繁重工作的 Python 程序——似乎在高性能、數(shù)值、桌面/服務(wù)器領(lǐng)域占據(jù)主導(dǎo)地位。雖然我不確定這個事實是否非常鼓舞人心，但我認為這是一個事實 ，而且這種情況將會持續(xù)很長時間。
引用鏈接
[1] 原文：A 100x speedup with unsafe Python: https://yosefk.com/blog/a-100x-speedup-with-unsafe-python.html[2] GitHub: https://github.com/yosefk/BlogCodeSamples/blob/main/numpy-perf.py
如果你正在尋找優(yōu)質(zhì)的Python文章和項目，我必須向你推薦??Python潮流周刊??！
它精選全網(wǎng)的優(yōu)秀文章、教程、開源項目、軟件工具、播客、視頻、熱門話題等豐富內(nèi)容，讓你緊跟技術(shù)最前沿，獲取最新的第一手學(xué)習(xí)資料！
歡迎點擊下方圖片，了解這份全世界知識密度最高、知識廣度最大的 Python 技術(shù)周刊。