Python 张量运算中的内存和时间

目标 我的目标是计算由下面的公式给出的张量。指数i，j，k，l从0到40，p，m，x从0到80

张量法这个总和只是收缩了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秒，我想重复多次

嵌套循环方法中的计时目标将时间减少10倍是令人满意的，但如果可以将时间减少100倍，那就足够了

你有没有解决内存问题或减少时间的想法？如何处理具有附加约束的此类求和

（备注：矩阵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

---

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解决方案的速度大约快了三个数量级，并且不应该占用超过必要的内存。

您能否展示您使用

tensordot

和嵌套循环的尝试，以便我们更好地了解问题？您也尝试过吗？比如：

t=tf.einsum（'pm，ijpx，klmx->ijkl'，a，b，b）

。不确定它对tensordot是否有任何影响，但可以尝试一下。您在TensorFlow或NumPy中也需要它吗？看起来您的代码是NumPy，但标记包括TensorFlow。我在NumPy工作，很抱歉标记有误导性。我把问题更新了一点。不，我没有尝试过tf.einsum方法。只是想澄清一下，einsum比嵌套循环长得多，因为它没有使用额外的约束。经过5分钟的等待，我结束了整个过程。非常感谢！！：）您的代码工作得非常出色。时间缩短了100倍，这是我的时间天堂。我刚刚编辑了线di，dj，dx，=B.shape为di，dj，dx，=（40,40,80,80），因为输出张量应该具有尺寸为40的所有尺寸。

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)