加速将24位二进制数据加载到16位numpy数组中

Question

我用下面的代码把24位的二进制 data 加载到一个16位的 numpy 数组里：

temp = numpy.zeros((len(data) / 3, 4), dtype='b')
temp[:, 1:] = numpy.frombuffer(data, dtype='b').reshape(-1, 3)
temp2 = temp.view('<i4').flatten() >> 16       # >> 16 because I need to divide by 2**16 to load my data into 16-bit array, needed for my (audio) application
output = temp2.astype('int16')

我想这可能有办法提高 速度效率，但是怎么做呢？

Answer 1

受到@amaurea的回答启发，这里有一个cython版本（我在原始代码中已经使用了cython，所以我会继续使用cython，而不是混合使用cython和fortran）:

import cython
import numpy as np
cimport numpy as np

def binary24_to_int16(char *data):
    cdef int i
    res = np.zeros(len(data)/3, np.int16)
    b = <char *>((<np.ndarray>res).data)
    for i in range(len(data)/3):
        b[2*i] = data[3*i+1]
        b[2*i+1] = data[3*i+2]
    return res            

速度提升了4倍哦 :)

Answer 2

看起来你这里说得有点绕。这样做不是能达到同样的效果吗？

output = np.frombuffer(data,'b').reshape(-1,3)[:,1:].flatten().view('i2')

这样做可以节省一些时间，因为不需要给临时数组填充零，省去了位移操作，也避免了一些不必要的数据移动。不过，我还没有实际测试过这个方法，预计节省的时间不会太多。

更新：我现在已经进行了测试。对于长度为1200万的数据，我的版本用时39毫秒，而你的版本用时80毫秒，速度大约快了两倍。虽然提升不算大，但也在预期之内，因为你的方法本身就已经很快了。

更新2：我应该提到我这里假设使用的是小端格式。不过，原问题中的代码也隐含地假设了小端格式，所以这并不是我新加的假设。

（如果是大端格式（数据和架构），你需要把1:替换成:-1。如果数据的字节顺序和CPU不同，你还需要反转字节的顺序（::-1）。）

更新3：为了获得更快的速度，我觉得你可能需要用到Python以外的东西。这个Fortran函数也使用了openMP，相比我的版本，它的速度提升超过两倍（所以比你的快四倍以上）。

subroutine f(a,b)
        implicit none
        integer*1, intent(in)  :: a(:)
        integer*1, intent(out) :: b(size(a)*2/3)
        integer :: i
        !$omp parallel do
        do i = 1, size(a)/3
                b(2*(i-1)+1) = a(3*(i-1)+2)
                b(2*(i-1)+2) = a(3*(i-1)+3)
        end do
        !$omp end parallel do
end subroutine

编译时使用FOPT="-fopenmp" f2py -c -m basj{,.f90} -lgomp。然后你可以在Python中导入并使用它：

import basj
def convert(data): return def mine2(data): return basj.f(np.frombuffer(data,'b')).view('i2')

你可以通过环境变量OMP_NUM_THREADS来控制使用的核心数量，但默认是使用所有可用的核心。