如何加速numpy數組屏蔽？

Question

我正在使用line profiler分析一段 Python 代碼的性能。

在代碼中，我有一個形狀為 (106906,) 和 dtype= int64的 numpy 數組tt 。 在探查器的幫助下，我發現mask[tt]=True下面的第二行很慢。 有沒有辦法加速它？ 如果這很重要，我正在使用 Python 3。

   mask = np.zeros(100000, dtype='bool')
   mask[tt] = True

Answer 1

您可以按照@orlevii 的建議使用 Numba：

from numba import njit
@njit
def f(mask,tt):
    mask[tt] = True
#Test:
mask = np.zeros(1000000, dtype='bool')
tt = np.random.randint(0,1000000,106906)
f(mask,tt)

一個簡單的%%timeit檢查表明您應該期望執行速度大約快 3 倍。

可以通過利用 GPU 實現進一步的加速。 如何使用 PyTorch 執行此操作的示例：

import torch
mask = torch.zeros(1000000).type(torch.cuda.FloatTensor)
tt = torch.randint(0,1000000,torch.Size([106906])).type(torch.cuda.LongTensor)
mask[tt] =  True

請注意，這里我們使用了一個torch.Tensor對象，它相當於numpy.ndarray中的 numpy.ndarray。 僅當您擁有帶有 CUDA 的 GPU（NVIDIA）時，代碼才會運行。 在 Tesla V100-SXM2 上期望 x30 加速 wrt 您的原始代碼。