Tăng tốc cho vòng lặp trong convolution cho mảng 3D numpy?

Question

Thực hiện convolution dọc theo Z vector của một mảng numpy 3d, sau đó hoạt động khác trên các kết quả, nhưng nó là chậm như nó được thực hiện ngay bây giờ. Là vòng lặp cho những gì đang làm chậm tôi xuống đây hoặc là sự hòa giải? Tôi đã cố gắng định hình lại một vector 1d và thực hiện các convolution trong 1 pass (như tôi đã làm trong Matlab), mà không có vòng lặp for, nhưng nó không cải thiện hiệu suất. Phiên bản Matlab của tôi nhanh hơn khoảng 50% so với bất kỳ thứ gì tôi có thể đưa ra bằng Python. Phần có liên quan của mã:

convolved=np.zeros((y_lines,x_lines,z_depth)) 
for i in range(0, y_lines): 
    for j in range(0, x_lines): 
     convolved[i,j,:]= fftconvolve(data[i,j,:], Gauss) #80% of time here 
     result[i,j,:]= other_calculations(convolved[i,j,:]) #20% of time here 

Có cách nào tốt hơn để làm điều này hơn vòng lặp không? Nghe nói về Cython nhưng tôi có kinh nghiệm hạn chế trong Python như của bây giờ, sẽ nhằm mục đích cho các giải pháp đơn giản nhất.

Answer 1

Chức năng fftconvolve bạn đang sử dụng có lẽ là từ SciPy. Nếu vậy, hãy lưu ý rằng nó có các mảng N-chiều. Vì vậy, một cách nhanh hơn để làm chập của bạn sẽ được để tạo ra hạt nhân 3d tương ứng với không làm gì trong x và y kích thước và làm một chập gaussian 1d trong z.

Một số mã và thời gian kết quả dưới đây. Trên máy tính của tôi và với một số dữ liệu đồ chơi, điều này dẫn đến sự tăng tốc 10 ×, như bạn có thể thấy:

import numpy as np 
from scipy.signal import fftconvolve 
from scipy.ndimage.filters import gaussian_filter 

# use scipy filtering functions designed to apply kernels to isolate a 1d gaussian kernel 
kernel_base = np.ones(shape=(5)) 
kernel_1d = gaussian_filter(kernel_base, sigma=1, mode='constant') 
kernel_1d = kernel_1d/np.sum(kernel_1d) 

# make the 3d kernel that does gaussian convolution in z axis only 
kernel_3d = np.zeros(shape=(1, 1, 5,)) 
kernel_3d[0, 0, :] = kernel_1d 

# generate random data 
data = np.random.random(size=(50, 50, 50)) 

# define a function for loop based convolution for easy timeit invocation 
def convolve_with_loops(data): 
    nx, ny, nz = data.shape 
    convolved=np.zeros((nx, ny, nz)) 
    for i in range(0, nx): 
     for j in range(0, ny): 
      convolved[i,j,:]= fftconvolve(data[i, j, :], kernel_1d, mode='same') 
    return convolved 

# compute the convolution two diff. ways: with loops (first) or as a 3d convolution (2nd) 
convolved = convolve_with_loops(data) 
convolved_2 = fftconvolve(data, kernel_3d, mode='same') 

# raise an error unless the two computations return equivalent results 
assert np.all(np.isclose(convolved, convolved_2)) 

# time the two routes of the computation 
%timeit convolved = convolve_with_loops(data) 
%timeit convolved_2 = fftconvolve(data, kernel_3d, mode='same') 

---

timeit kết quả:

10 loops, best of 3: 198 ms per loop 
100 loops, best of 3: 18.1 ms per loop 

Answer 2

Tôi nghĩ rằng bạn đã tìm thấy source code của hàm fftconvolve. Thông thường, đối với đầu vào thực, nó sử dụng các hàm numpy.fft.rfftn và .irfftn, tính toán các phép biến đổi N-chiều. Vì mục đích là để làm biến đổi nhiều 1-D, về cơ bản bạn có thể viết lại fftconvolve như thế này (giản thể):

from scipy.signal.signaltools import _next_regular 

def fftconvolve_1d(in1, in2): 

    outlen = in1.shape[-1] + in2.shape[-1] - 1 
    n = _next_regular(outlen) 

    tr1 = np.fft.rfft(in1, n) 
    tr2 = np.fft.rfft(in2, n) 
    out = np.fft.irfft(tr1 * tr2, n) 

    return out[..., :outlen].copy() 

Và tính toán kết quả mong muốn:

result = fftconvolve_1d(data, Gauss) 

này hoạt động vì numpy.fft.rfft và .irfft (chú ý việc thiếu n trong tên) biến đổi trên một trục của mảng đầu vào (trục cuối cùng theo mặc định). Đây là khoảng 40% nhanh hơn so với mã OP trên hệ thống của tôi.

---

Tăng tốc thêm có thể đạt được bằng cách sử dụng một back-end FFT khác.

Đối với một, các chức năng trong scipy.fftpack có vẻ hơi nhanh hơn tương đương Numpy của chúng. Tuy nhiên, định dạng đầu ra của các biến thể Scipy là khá khó xử (xem docs) và điều này làm cho nó khó khăn để làm phép nhân.

Một chương trình phụ trợ khác có thể là FFTW thông qua trình bao bọc pyFFTW. Nhược điểm là biến đổi được bắt đầu bằng một "giai đoạn lập kế hoạch" chậm và các đầu vào phải được căn chỉnh 16 byte để đạt được hiệu suất tốt nhất. Điều này được giải thích khá rõ trong pyFFTW tutorial. Mã kết quả có thể là ví dụ:

from scipy.signal.signaltools import _next_regular 
import pyfftw 
pyfftw.interfaces.cache.enable() # Cache for the "planning" 
pyfftw.interfaces.cache.set_keepalive_time(1.0) 

def fftwconvolve_1d(in1, in2): 

    outlen = in1.shape[-1] + in2.shape[-1] - 1 
    n = _next_regular(outlen) 

    tr1 = pyfftw.interfaces.numpy_fft.rfft(in1, n) 
    tr2 = pyfftw.interfaces.numpy_fft.rfft(in2, n) 

    sh = np.broadcast(tr1, tr2).shape 
    dt = np.common_type(tr1, tr2) 
    pr = pyfftw.n_byte_align_empty(sh, 16, dt) 
    np.multiply(tr1, tr2, out=pr) 
    out = pyfftw.interfaces.numpy_fft.irfft(pr, n) 

    return out[..., :outlen].copy() 

Với đầu vào liên kết và "lập kế hoạch" được lưu trong bộ nhớ cache, tôi thấy tốc độ gần gấp 3 lần so với mã trong OP.Liên kết bộ nhớ can be easily checked bằng cách xem địa chỉ bộ nhớ trong thuộc tính ctypes.data của mảng Numpy.