代码优化帮助（循环到张量操作）答案

【问题标题】：Code optimization help (loop to tensor operation)代码优化帮助（循环到张量操作）
【发布时间】：2020-04-17 00:16:09
【问题描述】：

我有一个很长的循环 2*40*40*40*40（大约 500 万次迭代），我觉得它没有必要，可以简化为张量操作。原始代码运行大约 34 秒，循环中有很多数值运算，我设法将它们减少到一个常数，使代码在 7.8 秒内运行。但是我仍然觉得如果我删除循环并用张量操作替换它们，它可以减少到

 for t in range(2):
    for i2 in range(M1):
        for j2 in range(M2):
            for i1 in range(M1):
                for j1 in range(M2):
                    U_out[i2][j2][t] += U_in[i1][j1][t] * constants[j1,i1,j2,i2,t]

更新：经过更多努力，我设法将时间减少到 0.1 秒：

for t in range(2):
  for i2 in range(M1):
     for j2 in range(M2):
       U_out[i2,j2,t] += np.sum(U_in[:,:,t] * constants[:,:,j2,i2,t])

感觉就像我错过了拼图的最后一部分。

【问题讨论】：

很高兴您能够对此进行优化。将 34 秒降低到 0.1 是一项了不起的成就，如果它真的有用，我建议只进一步优化它。如果您正在寻找更多输入，包含reproducible example 可能会很有用。通过扩展一些数组并利用爱因斯坦符号，您可以进一步将这个问题重新构建为跨几个维度的矩阵数学
docs.scipy.org/doc/numpy/reference/generated/numpy.einsum.html

标签： python numpy loops optimization tensor

【解决方案1】：

请调整此示例中的尺寸，我不确定我是否正确。所以我们有一个小的测试数组可以使用并比较结果：

import numpy as np

U_out = np.empty((4, 4, 2))

constants = np.random.random_sample((4, 4, 4, 4, 2))
U_in = np.random.random_sample((4, 4, 2))

for t in range(2):
    for i2 in range(4):
        for j2 in range(4):
            for i1 in range(4):
                for j1 in range(4):
                    U_out[i2][j2][t] += U_in[i1][j1][t] * constants[j1,i1,j2,i2,t]

U_out2 = U_out.copy()
print(U_out)

for t in range(2):
  for i2 in range(4):
     for j2 in range(4):
       U_out[i2,j2,t] += np.sum(U_in[:,:,t] * constants[:,:,j2,i2,t])

print(U_out - U_out2)

【讨论】：