如何在Pytorch中有效地计算张量？

Question

我有张量x和x.shape=(batch_size,10)，现在我要

 x[i][0] = x[i][0]*x[i][1]*...*x[i][9] for i in range(batch_size)

这是我的代码：

for i in range(batch_size):
    for k in range(1, 10):
        x[i][0] = x[i][0] * x[i][k]

但是当我在forward()中实现并调用loss.backward()时，反向传播的速度非常慢。为什么它很慢，并且有什么方法可以有效地实现它？

Answer 1

这很慢，因为您使用了两个for循环。

您可以使用.prod参见：https://pytorch.org/docs/stable/torch.html#torch.prod

就您而言，

x = torch.prod(x, dim=1)或x = x.prod(dim=1)

应该工作

Answer 2

使用两个循环计算乘积时，复杂度为n ^ 2。想象一下，在反向传播过程中这样做了很多次，您的代码变慢了。

向量运算可加快这些计算的速度。