张量运算中的内存和时间 python
Memory and time in tensor operations python
目标
我的目标是计算您可以在下面看到的公式给出的张量。指数 i, j, k, l 运行 从 0 到 40 和 p, m, x 从 0 到 80.
Tensordot approach这个求和只是收缩了6个指数的巨大张量。我尝试通过允许进行此类计算的张量点来完成,但即使我先做一个张量点,然后再做另一个,我的问题还是内存。 (我在 colab 工作,所以我有 12GB RAM 可用)
嵌套循环方法但是有一些额外的对称性支配 B 矩阵,即 B{ijpx} 的唯一非零元素是 i+j= p+x。因此,我能够将 p 和 m 写成 x 的函数(p=i+j-x,m=k+l-x),然后我做了 5 个循环,即 i、j、k、l、x,但另一方面时间是问题,因为计算需要 136 秒,我想重复多次。
嵌套循环方法中的计时目标将时间减少十倍是令人满意的,但如果有可能将其减少一百倍,那将是绰绰有余。
对于解决内存问题或减少时间,您有什么想法吗?您如何处理带有额外约束的此类求和?
(注:矩阵A是对称的,我至今没有用到这个事实,没有更多的对称性了。)
这里是嵌套循环的代码:
for i in range (0,40):
for j in range (0,40):
for k in range (0,40):
for l in range (0,40):
Sum=0
for x in range (0,80):
p=i+j-x
m=k+l-x
if p>=0 and p<80 and m>=0 and m<80:
Sum += A[p,m]*B[i,j,p,x]*B[k,l,m,x]
T[i,j,k,l]= Sum
以及张量点方法的代码:
P=np.tensordot(A,B,axes=((0),(2)))
T=np.tensordot(P,B,axes=((0,3),(2,3)))
Numba 可能是您最好的选择。我根据您的代码将此功能放在一起。我稍微改变了它以避免一些不必要的迭代和 if
块:
import numpy as np
import numba as nb
@nb.njit(parallel=True)
def my_formula_nb(A, B):
di, dj, dx, _ = B.shape
T = np.zeros((di, dj, di, dj), dtype=A.dtype)
for i in nb.prange (di):
for j in nb.prange (dj):
for k in nb.prange (di):
for l in nb.prange (dj):
sum = 0
x_start = max(0, i + j - dx + 1, k + l - dx + 1)
x_end = min(dx, i + j + 1, k + l + 1)
for x in range(x_start, x_end):
p = i + j - x
m = k + l - x
sum += A[p, m] * B[i, j, p, x] * B[k, l, m, x]
T[i, j, k, l] = sum
return T
让我们看看实际效果:
import numpy as np
def make_problem(di, dj, dx):
a = np.random.rand(dx, dx)
a = a + a.T
b = np.random.rand(di, dj, dx, dx)
b_ind = np.indices(b.shape)
b_mask = b_ind[0] + b_ind[1] != b_ind[2] + b_ind[3]
b[b_mask] = 0
return a, b
# Generate a problem
np.random.seed(100)
a, b = make_problem(15, 20, 25)
# Solve with Numba function
t1 = my_formula_nb(a, b)
# Solve with einsum
t2 = np.einsum('pm,ijpx,klmx->ijkl', a, b, b)
# Check result
print(np.allclose(t1, t2))
# True
# Benchmark (IPython)
%timeit np.einsum('pm,ijpx,klmx->ijkl', a, b, b)
# 4.5 s ± 39.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit my_formula_nb(a, b)
# 6.06 ms ± 20.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
如您所见,Numba 解决方案的速度大约快了三个数量级,而且它不会占用比必要更多的内存。
目标 我的目标是计算您可以在下面看到的公式给出的张量。指数 i, j, k, l 运行 从 0 到 40 和 p, m, x 从 0 到 80.
Tensordot approach这个求和只是收缩了6个指数的巨大张量。我尝试通过允许进行此类计算的张量点来完成,但即使我先做一个张量点,然后再做另一个,我的问题还是内存。 (我在 colab 工作,所以我有 12GB RAM 可用)
嵌套循环方法但是有一些额外的对称性支配 B 矩阵,即 B{ijpx} 的唯一非零元素是 i+j= p+x。因此,我能够将 p 和 m 写成 x 的函数(p=i+j-x,m=k+l-x),然后我做了 5 个循环,即 i、j、k、l、x,但另一方面时间是问题,因为计算需要 136 秒,我想重复多次。
嵌套循环方法中的计时目标将时间减少十倍是令人满意的,但如果有可能将其减少一百倍,那将是绰绰有余。
对于解决内存问题或减少时间,您有什么想法吗?您如何处理带有额外约束的此类求和?
(注:矩阵A是对称的,我至今没有用到这个事实,没有更多的对称性了。)
这里是嵌套循环的代码:
for i in range (0,40):
for j in range (0,40):
for k in range (0,40):
for l in range (0,40):
Sum=0
for x in range (0,80):
p=i+j-x
m=k+l-x
if p>=0 and p<80 and m>=0 and m<80:
Sum += A[p,m]*B[i,j,p,x]*B[k,l,m,x]
T[i,j,k,l]= Sum
以及张量点方法的代码:
P=np.tensordot(A,B,axes=((0),(2)))
T=np.tensordot(P,B,axes=((0,3),(2,3)))
Numba 可能是您最好的选择。我根据您的代码将此功能放在一起。我稍微改变了它以避免一些不必要的迭代和 if
块:
import numpy as np
import numba as nb
@nb.njit(parallel=True)
def my_formula_nb(A, B):
di, dj, dx, _ = B.shape
T = np.zeros((di, dj, di, dj), dtype=A.dtype)
for i in nb.prange (di):
for j in nb.prange (dj):
for k in nb.prange (di):
for l in nb.prange (dj):
sum = 0
x_start = max(0, i + j - dx + 1, k + l - dx + 1)
x_end = min(dx, i + j + 1, k + l + 1)
for x in range(x_start, x_end):
p = i + j - x
m = k + l - x
sum += A[p, m] * B[i, j, p, x] * B[k, l, m, x]
T[i, j, k, l] = sum
return T
让我们看看实际效果:
import numpy as np
def make_problem(di, dj, dx):
a = np.random.rand(dx, dx)
a = a + a.T
b = np.random.rand(di, dj, dx, dx)
b_ind = np.indices(b.shape)
b_mask = b_ind[0] + b_ind[1] != b_ind[2] + b_ind[3]
b[b_mask] = 0
return a, b
# Generate a problem
np.random.seed(100)
a, b = make_problem(15, 20, 25)
# Solve with Numba function
t1 = my_formula_nb(a, b)
# Solve with einsum
t2 = np.einsum('pm,ijpx,klmx->ijkl', a, b, b)
# Check result
print(np.allclose(t1, t2))
# True
# Benchmark (IPython)
%timeit np.einsum('pm,ijpx,klmx->ijkl', a, b, b)
# 4.5 s ± 39.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit my_formula_nb(a, b)
# 6.06 ms ± 20.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
如您所见,Numba 解决方案的速度大约快了三个数量级,而且它不会占用比必要更多的内存。