首頁猿問蟒蛇中FFT的循環(huán)加速（使用&ld...

蟒蛇中FFT的循環(huán)加速（使用“np.einsum”）

Python

寶慕林4294392 2022-09-20 17:46:21

問題：我想加速我的python循環(huán)，其中包含很多產(chǎn)品和總結(jié)，但我也對任何其他解決方案持開放態(tài)度。np.einsum我的函數(shù)采用形狀為（n，n，3）的向量配置 S（我的情況：n=72），并對 N*N 點(diǎn)的相關(guān)函數(shù)進(jìn)行傅里葉變換。相關(guān)函數(shù)定義為每個向量與其他向量的乘積。這乘以向量位置乘以 kx 和 ky 值的余弦函數(shù)。每個位置在最終求和得到k空間中的一個點(diǎn)：i,jp,mdef spin_spin(S,N): n= len(S) conf = np.reshape(S,(n**2,3)) chi = np.zeros((N,N)) kx = np.linspace(-5*np.pi/3,5*np.pi/3,N) ky = np.linspace(-3*np.pi/np.sqrt(3),3*np.pi/np.sqrt(3),N) x=np.reshape(triangular(n)[0],(n**2)) y=np.reshape(triangular(n)[1],(n**2)) for p in range(N): for m in range(N): for i in range(n**2): for j in range(n**2): chi[p,m] += 2/(n**2)*np.dot(conf[i],conf[j])*np.cos(kx[p]*(x[i]-x[j])+ ky[m]*(y[i]-y[j])) return(chi,kx,ky)我的問題是我需要大約100 * 100個點(diǎn)，這些點(diǎn)由kx * ky表示，并且循環(huán)需要幾個小時才能完成具有72 * 72個向量的格的工作。計算次數(shù)：72 * 72 * 72 * 72 * 100 * 100 我不能使用內(nèi)置的FFT，因為我的三角形網(wǎng)格，所以我需要一些其他選項來降低這里的計算成本。numpy我的想法：首先，我意識到將配置重新塑造為向量列表而不是矩陣可以降低計算成本。此外，我使用了numba包，這也降低了成本，但它仍然太慢了。我發(fā)現(xiàn)計算這些對象的一個好方法是函數(shù)。計算每個向量與每個向量的乘積是通過以下方式完成的：np.einsumnp.einsum('ij,kj -> ik',np.reshape(S,(72**2,3)),np.reshape(S,(72**2,3)))棘手的部分是計算 .在這里，我想將產(chǎn)品與向量的位置（例如）的形狀列表（100，1）連接起來。特別是我真的不知道如何用來實(shí)現(xiàn)x方向和y方向的距離。np.cosnp.shape(x)=(72**2,1)np.einsum要重現(xiàn)代碼（可能不需要這樣）：首先，您需要一個矢量配置。你可以簡單地用它來做，或者你用隨機(jī)向量作為例子：np.ones((72,72,3)def spherical_to_cartesian(r, theta, phi): '''Convert spherical coordinates (physics convention) to cartesian coordinates''' sin_theta = np.sin(theta) x = r * sin_theta * np.cos(phi) y = r * sin_theta * np.sin(phi) z = r * np.cos(theta) return x, y, z # return a tupledef random_directions(n, r): '''Return ``n`` 3-vectors in random directions with radius ``r``''' out = np.empty(shape=(n,3), dtype=np.float64) for i in range(n): # Pick directions randomly in solid angle phi = random.uniform(0, 2*np.pi) theta = np.arccos(random.uniform(-1, 1)) # unpack a tuple x, y, z = spherical_to_cartesian(r, theta, phi) out[i] = x, y, z

查看完整描述

2 回答

翻翻過去那場雪

TA貢獻(xiàn)2065條經(jīng)驗獲得超14個贊

優(yōu)化的努姆巴實(shí)施

代碼中的主要問題是使用極小的數(shù)據(jù)重復(fù)調(diào)用外部 BLAS 函數(shù)。在此代碼中，只計算一次它們更有意義，但是如果您必須在循環(huán)中執(zhí)行此計算，請編寫Numba實(shí)現(xiàn)。例np.dot

優(yōu)化的功能（暴力破解）

import numpy as np

import numba as nb

@nb.njit(fastmath=True,error_model="numpy",parallel=True)

def spin_spin(S,N):

n= len(S)

conf = np.reshape(S,(n**2,3))

chi = np.zeros((N,N))

kx = np.linspace(-5*np.pi/3,5*np.pi/3,N).astype(np.float32)

ky = np.linspace(-3*np.pi/np.sqrt(3),3*np.pi/np.sqrt(3),N).astype(np.float32)

x=np.reshape(triangular(n)[0],(n**2)).astype(np.float32)

y=np.reshape(triangular(n)[1],(n**2)).astype(np.float32)

#precalc some values

fact=nb.float32(2/(n**2))

conf_dot=np.dot(conf,conf.T).astype(np.float32)

for p in nb.prange(N):

for m in range(N):

#accumulating on a scalar is often beneficial

acc=nb.float32(0)

for i in range(n**2):

for j in range(n**2):

acc+= conf_dot[i,j]*np.cos(kx[p]*(x[i]-x[j])+ ky[m]*(y[i]-y[j]))

chi[p,m]=fact*acc

return(chi,kx,ky)

優(yōu)化功能（刪除冗余計算）

有很多冗余的計算。這是有關(guān)如何刪除它們的示例。這也是一個以雙精度進(jìn)行計算的版本。

@nb.njit()

def precalc(S):

#There may not be all redundancies removed

n= len(S)

conf = np.reshape(S,(n**2,3))

conf_dot=np.dot(conf,conf.T)

x=np.reshape(triangular(n)[0],(n**2))

y=np.reshape(triangular(n)[1],(n**2))

x_s=set()

y_s=set()

for i in range(n**2):

for j in range(n**2):

x_s.add((x[i]-x[j]))

y_s.add((y[i]-y[j]))

x_arr=np.sort(np.array(list(x_s)))

y_arr=np.sort(np.array(list(y_s)))

conf_dot_sel=np.zeros((x_arr.shape[0],y_arr.shape[0]))

for i in range(n**2):

for j in range(n**2):

ii=np.searchsorted(x_arr,x[i]-x[j])

jj=np.searchsorted(y_arr,y[i]-y[j])

conf_dot_sel[ii,jj]+=conf_dot[i,j]

return x_arr,y_arr,conf_dot_sel

@nb.njit(fastmath=True,error_model="numpy",parallel=True)

def spin_spin_opt_2(S,N):

chi = np.empty((N,N))

n= len(S)

kx = np.linspace(-5*np.pi/3,5*np.pi/3,N)

ky = np.linspace(-3*np.pi/np.sqrt(3),3*np.pi/np.sqrt(3),N)

x_arr,y_arr,conf_dot_sel=precalc(S)

fact=2/(n**2)

for p in nb.prange(N):

for m in range(N):

acc=nb.float32(0)

for i in range(x_arr.shape[0]):

for j in range(y_arr.shape[0]):

acc+= fact*conf_dot_sel[i,j]*np.cos(kx[p]*x_arr[i]+ ky[m]*y_arr[j])

chi[p,m]=acc

return(chi,kx,ky)

@nb.njit()

def precalc(S):

#There may not be all redundancies removed

n= len(S)

conf = np.reshape(S,(n**2,3))

conf_dot=np.dot(conf,conf.T)

x=np.reshape(triangular(n)[0],(n**2))

y=np.reshape(triangular(n)[1],(n**2))

x_s=set()

y_s=set()

for i in range(n**2):

for j in range(n**2):

x_s.add((x[i]-x[j]))

y_s.add((y[i]-y[j]))

x_arr=np.sort(np.array(list(x_s)))

y_arr=np.sort(np.array(list(y_s)))

conf_dot_sel=np.zeros((x_arr.shape[0],y_arr.shape[0]))

for i in range(n**2):

for j in range(n**2):

ii=np.searchsorted(x_arr,x[i]-x[j])

jj=np.searchsorted(y_arr,y[i]-y[j])

conf_dot_sel[ii,jj]+=conf_dot[i,j]

return x_arr,y_arr,conf_dot_sel

@nb.njit(fastmath=True,error_model="numpy",parallel=True)

def spin_spin_opt_2(S,N):

chi = np.empty((N,N))

n= len(S)

kx = np.linspace(-5*np.pi/3,5*np.pi/3,N)

ky = np.linspace(-3*np.pi/np.sqrt(3),3*np.pi/np.sqrt(3),N)

x_arr,y_arr,conf_dot_sel=precalc(S)

fact=2/(n**2)

for p in nb.prange(N):

for m in range(N):

acc=nb.float32(0)

for i in range(x_arr.shape[0]):

for j in range(y_arr.shape[0]):

acc+= fact*conf_dot_sel[i,j]*np.cos(kx[p]*x_arr[i]+ ky[m]*y_arr[j])

chi[p,m]=acc

return(chi,kx,ky)

計時

#brute-force

%timeit res=spin_spin(S,100)

#48 s ± 671 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

#new version

%timeit res_2=spin_spin_opt_2(S,100)

#5.33 s ± 59.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit res_2=spin_spin_opt_2(S,1000)

#1min 23s ± 2.43 s per loop (mean ± std. dev. of 7 runs, 1 loop each)

編輯（安全監(jiān)控系統(tǒng)檢查）

import numba as nb

import numpy as np

@nb.njit(fastmath=True)

def foo(n):

x = np.empty(n*8, dtype=np.float64)

ret = np.empty_like(x)

for i in range(ret.size):

ret[i] += np.cos(x[i])

return ret

foo(1000)

if 'intel_svmlcc' in foo.inspect_llvm(foo.signatures[0]):

print("found")

else:

print("not found")

#found

如果有閱讀此鏈接。它應(yīng)該可以在Linux和Windows上運(yùn)行，但我還沒有在macOS上測試過它。not found

反對回復(fù) 2022-09-20

胡子哥哥

TA貢獻(xiàn)1825條經(jīng)驗獲得超6個贊

這是加快速度的一種方法。我沒有開始使用np.einsum，因為對你的循環(huán)進(jìn)行一些調(diào)整就足夠了。

減慢代碼速度的主要因素是同一事物的冗余重新計算。此處的嵌套循環(huán)是犯罪者：

for p in range(N):

for m in range(N):

for i in range(n**2):

for j in range(n**2):

chi[p,m] += 2/(n**2)*np.dot(conf[i],conf[j])*np.cos(kx[p]*(x[i]-x[j])+ ky[m]*(y[i]-y[j]))

它包含大量冗余，多次重新計算向量運(yùn)算。

考慮 np.dot（...）：此計算完全獨(dú)立于點(diǎn) kx 和 ky。但只有點(diǎn) kx 和 ky 需要用 m 和 n 進(jìn)行索引。因此，您可以對所有i和j運(yùn)行一次點(diǎn)積，并保存結(jié)果，而不是為每個m，n重新計算（這將是10，000次！

在類似的方法中，不需要在晶格中的每個點(diǎn)重新計算兩者之間的向量差異。在每個點(diǎn)上，您都會計算每個矢量距離，而只需要計算一次矢量距離，然后只需將此結(jié)果乘以每個晶格點(diǎn)即可。

因此，在修復(fù)了循環(huán)并使用索引（i，j）作為鍵的字典來存儲所有值之后，您只需在i，j上的循環(huán)中查找相關(guān)值即可。這是我的代碼：

def spin_spin(S, N):

n = len(S)

conf = np.reshape(S,(n**2, 3))

chi = np.zeros((N, N))

kx = np.linspace(-5*np.pi/3, 5*np.pi/3, N)

ky = np.linspace(-3*np.pi/np.sqrt(3), 3*np.pi/np.sqrt(3), N)

# Minor point; no need to use triangular twice

x, y = triangular(n)

x, y = np.reshape(x,(n**2)), np.reshape(y,(n**2))

# Build a look-up for all the dot products to save calculating them many times

dot_prods = dict()

x_diffs, y_diffs = dict(), dict()

for i, j in itertools.product(range(n**2), range(n**2)):

dot_prods[(i, j)] = np.dot(conf[i], conf[j])

x_diffs[(i, j)], y_diffs[(i, j)] = x[i] - x[j], y[i] - y[j]

# Minor point; improve syntax by converting nested for loops to one line

for p, m in itertools.product(range(N), range(N)):

for i, j in itertools.product(range(n**2), range(n**2)):

# All vector operations are replaced by look ups to the dictionaries defined above

chi[p, m] += 2/(n**2)*dot_prods[(i, j)]*np.cos(kx[p]*(x_diffs[(i, j)]) + ky[m]*(y_diffs[(i, j)]))

return(chi, kx, ky)

我現(xiàn)在正在一臺像樣的機(jī)器上使用您提供的尺寸運(yùn)行此內(nèi)容，并且i，j上的循環(huán)將在兩分鐘內(nèi)完成。這只需要發(fā)生一次;那么它只是一個在m，n上的循環(huán)。每一個大約需要90秒，所以仍然有2-3個小時的運(yùn)行時間。我歡迎任何關(guān)于如何優(yōu)化cos計算以加快速度的建議！

我擊中了優(yōu)化的唾手可得的果實(shí)，但為了給人一種速度感，i，j的循環(huán)需要2分鐘，這樣它運(yùn)行的次數(shù)減少了9，999次！

反對回復(fù) 2022-09-20