Question

我正在学习cython。我写了一个numpy代码：

from numpy import *

def set_onsite(n):
    a=linspace(0,n,n+1)
    onsite=zeros([n+1,n+1],float)
    for i in range(0,n+1):
        onsite[i,i]=a[i]*a[i]
    return onsite

我使用％timeit来计算时间

%timeit ex5.set_onsite(10000)
10 loops, best of 3: 24.2 ms per loop

所以我写了一个这样的cython版本：

import numpy as np
cimport numpy as np
cimport cython
import cython

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)

def set_onsite(np.int_t n):
    cdef np.ndarray[double,ndim=1,mode='c'] a=np.linspace(0,n,n+1)
    cdef np.ndarray[double,ndim=2,mode='c'] onsite=np.empty(n+1,n+1)
    cdef np.int_t i

    for i in range(0,n+1):
        onsite[i,i]=a[i]*a[i]
    return onsite

这次的结果是：

%timeit ex5.set_onsite(10000)
100 loops, best of 3: 18.1 ms per loop

结果几乎相同。我不满足于这个结果，所以我想知道是否有任何方法可以使代码比每个循环的18.1 ms更快？

Answer 1

Cython可能不是加速计算的正确工具。

首先要记住的是，矢量化可以为你带来巨大的速度提升。我们的想法是将显式for循环转换为整个向量上的操作。

此函数执行相同操作，但避免使用for循环：

def set_onsite_vec(n):
    a = np.linspace(0,n,n+1)
    diag = a*a
    return np.diag(diag)

%timeit set_onsite(100) - ＆gt; 10000 loops, best of 3: 39 µs per loop %timeit set_onsite_vec(100) - ＆gt; `10000循环，最佳3：每循环16.3μs``

这会给你带来很大的提升，但你会注意到，对于较大的阵列，这种增益会降低。

问题是你正在创建一个巨大的矩阵（O(n^2)条目） - 必须分配和初始化这个矩阵，这是瓶颈。

请注意，矩阵在对角线上只有非零元素。您可以使用该信息并创建稀疏矩阵。

from scipy.sparse import diags

def set_onsite_sparse(n):
    a = np.linspace(0,n,n+1)
    diag = a*a
    return diags([diag], [0])

%timeit set_onsite_sparse(10000) - ＆gt; 1000 loops, best of 3: 119 µs per loop

不创建密集矩阵可以大大提高速度。

当然，如果在转换回密集矩阵之前进行其他稀疏操作，则只会使用该改进。

简而言之：使用Cython可能无法显着改善，但使用稀疏线性代数可以为您带来巨大的性能提升。但这取决于您希望如何在程序中使用此功能。

加速我的numpy代码

1 个答案: