如何使用torch提供的函数优化5层循环？

Question

x是形状为(16，10，4，25，53)的张量，y的大小与x相同。

mean的形状为(25，53)，jc和ac的大小均为(16，10，4)。

如何使用torch函数优化以下表达式？

for k in range(x.size()[0]):
    for s in range(x.size()[1]):
        for u in range(x.size()[2]):
            for i in range(x.size()[3]):
                for j in range(x.size()[4]):
                    num1 += (x[k][s][u][i][j] - mean[i][j] - jc[k][s][u]) * (y[k][s][u][i][j] - mean[i][j] - ac[k][s][u])
                    num2 += (y[k][s][u][i][j] - mean[i][j] - jc[k][s][u]) ** 2
                    num3 += (y[k][s][u][i][j] - mean[i][j] - ac[k][s][u]) ** 2

Answer 1

我认为你看到的是，沿着单元素维度的张量。，broadcasting。

首先，您需要维度的数量相同，因此如果mean的形状为(25,53)，那么mean[None, None, None, ...]的形状为(1, 1, 1, 25, 53) -您没有更改底层数据中的任何内容，但维度的数量现在是5而不是只有2，并且这些单一维度可以广播到x和y的相应维度。

使用广播的优化代码将如下所示：

num1 = ((x - mean[None, None, None, ...] - jc[..., None, None]) * (y - mean[None, None, None, ...] - ac[..., None, None])).sum()
num2 = ((y - mean[None, None, None, ...] - jc[..., None, None]) ** 2).sum()  # shouldn't it be x here?
num3 = ((y - mean[None, None, None, ...] - ac[..., None, None]) ** 2).sum()