【问题标题】:Fast way to bin a 2D array in python在 python 中对二维数组进行 bin 的快速方法
【发布时间】:2020-04-20 15:02:26
【问题描述】:

我正在寻找一种对二维数组进行分箱的快速方法。 我在这个网站上看到了这个话题,bin a 2d array in numpy,发现下面用numpy的解决方案可以快速处理数据。

编辑: 分箱是这样的, 如果有的话,

array([[  0,   0,   0,   0,   0,   0],
       [  0, 144,   0,   0,   0,   0],
       [  0,   0,   0, 144,   3,   0],
       [109, 112, 116, 121,  40,  91]])

binning 的输出将是

array([[144,   0,   0],
       [221, 381, 134]])

如您所见,在这种情况下,输出数组的每个元素都是原始数组中 2x2 数组的总和值。在我的情况下,这些垃圾箱大约为 50x50。

如果 a 的形状为 m, n,则 reshape 的形式应为 a.reshape(m_bins, m // m_bins, n_bins, n // n_bins)

但是因为我必须处理一个大数组(超过 1k x 1k),所以这在我的 PC 上需要几十毫秒。有什么方法可以更快地完成这项工作,例如在 Cython 中使用 C?

【问题讨论】:

  • 您能否提供一些示例输入和输出数据以及您正在使用的代码?
  • 内核 2x2 和步幅 2x2 的有效卷积应该可以快速工作,但可能有更好的方法。
  • 垃圾箱大约是 50x50

标签: python numpy cython


【解决方案1】:

这可能令人惊讶,但总结矩阵中的一些值并不是一件容易的事。我的Thisthis 的答案给出了一些见解,所以我不打算重复太多细节,但为了获得最佳性能,必须在现代 CPU 上利用缓存和 SIMD/管道性质的浮动操作最好的办法。

Numpy 获得了上述所有权利,并且很难被击败。这不是不可能的,但是一个人必须非常熟悉低级优化才能成功——而且人们不应该期望有太多的改进。幼稚的实现根本无法击败 Numpy。

这是我使用 cython 和 numba 的尝试,所有的加速都来自并行化。

让我们从基线算法开始:

def bin2d(a,K):
    m_bins = a.shape[0]//K
    n_bins = a.shape[1]//K
    return a.reshape(m_bins, K, n_bins, K).sum(3).sum(1)

并衡量它的性能:

import numpy as np
N,K=2000,50
a=np.arange(N*N, dtype=np.float64).reshape(N,N)
%timeit bin2d(a,K)
# 2.76 ms ± 107 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

赛通:

这是我使用 Cython 的实现,它使用 OpenMP 进行并行化。为了简单起见,我在原地执行求和,因此传递的数组将被更改(numpy 的版本不是这种情况):

%%cython -a  -c=/openmp --link-args=/openmp -c=/arch:AVX512
import numpy as np
import cython
from cython.parallel import prange
cdef extern from *:
    """
    //assumes direct memory and row-major-order (i.e. rows are continuous)
    double calc_bin(double *ptr, int N, int y_offset){
       double *row = ptr;
       for(int y=1;y<N;y++){
           row+=y_offset; //next row
           //there is no dependencies, so the summation can be done in parallel
           for(int x=0;x<N;x++){
               ptr[x]+=row[x];
           }
       }
       double res=0.0;
       for(int x=0;x<N;x++){
           //could be made slightly faster (summation is not in parallel), but is it needed?
           res+=ptr[x];
       }
       return res;
    }
    """
    double calc_bin(double *ptr, int N, int y_offset) nogil


@cython.boundscheck(False)
@cython.wraparound(False)
def cy_bin2d_parallel(double[:, ::1] a, int K):
    cdef int y_offset = a.shape[0]
    cdef int m_bins = a.shape[0]//K
    cdef int n_bins = a.shape[1]//K
    cdef double[:,:] res = np.empty((m_bins, n_bins), dtype=np.float64)
    cdef int i,j,k
    for k in prange(m_bins*n_bins, nogil=True):
        i = k//m_bins
        j = k%m_bins
        res[i,j] = calc_bin(&a[i*K, j*K], K, y_offset)
    return res.base

现在(验证结果相同后):

a=np.arange(N*N, dtype=np.float64).reshape(N,N)
%timeit cy_bin2d_parallel(a,K)
# 1.51 ms ± 25.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
# without parallelization: 2.93 ms

大约快 30%。使用@max9111 提出的-c=/arch:AVX512(否则只能使用/Ox 和MSVC 编译),使单线程版本的速度提高了约20%,但并行版本的速度仅提高了约5%。

Numba:

使用 numba 编译的算法相同(由于 clang-compiler 的性能更好,它通常可以击败 cython)——但结果比 cython 稍慢,但比 numpy 快 20%:

import numba as nb
@nb.njit(parallel=True)
def nb_bin2d_parallel(a, K):
    m_bins = a.shape[0]//K
    n_bins = a.shape[1]//K
    res = np.zeros((m_bins, n_bins), dtype=np.float64)

    for k in nb.prange(m_bins*n_bins):
        i = k//m_bins
        j = k%m_bins
        for y in range(i*K+1, (i+1)*K):
            for x in range(j*K, (j+1)*K):
                a[i*K, x] += a[y,x]
        s=0.0
        for x in range(j*K, (j+1)*K):
            s+=a[i*K, x]
        res[i,j] = s
    return res

现在:

a=np.arange(N*N, dtype=np.float64).reshape(N,N)
%timeit nb_bin2d_parallel(a,K)
# 1.98 ms ± 162 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
# without parallelization: 5.8 ms

简而言之:我想有可能击败上述方法,但不再有免费的午餐,因为 numpy 做得很好。最大的潜力可能是并行化,但由于问题的内存带宽限制性质而受到限制(并且可能应该像我一样使用更智能的策略 - 对于 50*50 的求和,一个可能的仍然可以看到开销创建/管理线程)。


Cython 有另一种更快的尝试(至少对于大约 2000 的大小),它不是对 50 个元素的小部分执行求和,而是对整行执行求和,从而减少开销(但一旦行大于 1-2k):

%%cython -a --verbose -c=/openmp --link-args=/openmp -c=/arch:AVX512
import numpy as np
import cython
from cython.parallel import prange
cdef extern from *:
    """
    void calc_bin_row(double *ptr, int N, int y_offset, double* out){
       double *row = ptr;
       for(int y=1;y<N;y++){
           row+=y_offset; //next row
           for(int x=0;x<y_offset;x++){
               ptr[x]+=row[x];
           }
       }
       double res=0.0;
       int i=0;
       int k=0;
       for(int x=0;x<y_offset;x++){//could be made faster, but is it needed?
           res+=ptr[x];
           k++;
           if(k==N){
              k=0;
              out[i]=res;
              i++;
              res=0.0;
           }
       }
    }
    """
    void calc_bin_row(double *ptr, int N, int y_offset, double* out) nogil


@cython.boundscheck(False)
@cython.wraparound(False)
def cy_bin2d_parallel_rowise(double[:, ::1] a, int K):
    cdef int y_offset = a.shape[1]
    cdef int m_bins = a.shape[0]//K
    cdef int n_bins = a.shape[1]//K
    cdef double[:,:] res = np.empty((m_bins, n_bins), dtype=np.float64)
    cdef int i,j,k
    for k in prange(0, y_offset, K, nogil=True):
        calc_bin_row(&a[k, 0], K, y_offset, &res[k//K, 0])
    return res.base

单线程 (2ms) 已经更快了,并行化速度提高了大约 60%(每个循环 1.27 ms ± 50.8 µs(平均值 ± 标准偏差,7 次运行,每次 100 个循环))。

【讨论】:

  • 如果您像在 C 代码中那样在第二个循环中求和一个标量,Numba 版本应该会快一点(并将结果写入数组 res 并分配 res使用 np.empty) 此外,如果函数尚未完全内存瓶颈,则 fastmath(Numba 和 Cython)、march=native(默认在 Numba 中,但在 Cython 中没有)可能会有所帮助。
  • @max9111,谢谢,在我看来,期望写入全局内存以及写入寄存器/局部变量会得到优化,这可能是天真的。这种变化带来了大约5%。然而,使用 fastmath 没有任何效果(正如所料,因为唯一的潜力是最后一个求和循环)。使用 /arch:AVX512(它是 MSVC)为 Cython 带来了进一步的加速。
  • 找到了一种更快的方法(仅在 Numba 中测试过),但也许您也可以在 Cython/C 中获得一些性能。我将添加一个“评论”答案。
【解决方案2】:

这基本上是对@ead answer的评论。

通常,在对齐求和上您能做的最好的事情就是尽可能多地使用标量(映射到寄存器)。此函数也不会修改输入数组,这可能是不需要的。

代码和时间

@nb.njit(parallel=True,fastmath=True,cache=True)
def nb_bin2d_parallel_2(a, K):
    #There is no bounds-checking, make sure that the dimensions are OK
    assert a.shape[0]%K==0
    assert a.shape[1]%K==0

    m_bins = a.shape[0]//K
    n_bins = a.shape[1]//K
    #Works for all datatypes, but overflow especially in small integer types
    #may occur
    res = np.zeros((m_bins, n_bins), dtype=a.dtype)

    for i in nb.prange(res.shape[0]):
        for ii in range(i*K,(i+1)*K):
            for j in range(res.shape[1]):
                TMP=res[i,j]
                for jj in range(j*K,(j+1)*K):
                    TMP+=a[ii,jj]
                res[i,j]=TMP
    return res

N,K=2000,50
a=np.arange(N*N, dtype=np.float64).reshape(N,N)

#warmup (Numba compilation is on the first call)
res_1=nb_bin2d_parallel(a, K)
res_2=cy_bin2d_parallel(a,K)
res_3=bin2d(a,K)
res_4=nb_bin2d_parallel_2(a, K)

%timeit bin2d(a,K)
#2.51 ms ± 25.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit nb_bin2d_parallel(a, K)
#1.33 ms ± 33.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit nb_bin2d_parallel_2(a, K)
#1.05 ms ± 8.96 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit cy_bin2d_parallel(a,K) #arch:AVX2
#996 µs ± 7.94 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

N,K=4000,50
a=np.arange(N*N, dtype=np.float64).reshape(N,N)

%timeit bin2d(a,K)
#10.8 ms ± 56.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit nb_bin2d_parallel(a, K)
#5.13 ms ± 46.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit nb_bin2d_parallel_2(a, K)
#3.99 ms ± 31.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit cy_bin2d_parallel(a,K) #arch:AVX2
#4.31 ms ± 168 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

【讨论】:

  • 没错,它比我的(旧)实现更舒服。我想有时我会过多地帮助优化器——他们只是比我好。我还发布了另一种方法,它比您的版本稍快 - 可能有办法进一步改进它,但必须真正了解瓶颈是什么。
猜你喜欢
  • 2020-12-04
  • 1970-01-01
  • 1970-01-01
  • 2016-07-03
  • 2020-11-01
  • 1970-01-01
  • 1970-01-01
  • 2017-09-14
  • 2020-05-23
相关资源
最近更新 更多