如何有效地计算 Pytorch 中的张量？答案

【问题标题】：How can I compute the tensor in Pytorch efficiently?如何有效地计算 Pytorch 中的张量？
【发布时间】：2019-05-10 23:59:10
【问题描述】：

我有张量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() 时，反向传播的速度非常慢。为什么它很慢，有什么方法可以有效地实现它？

【问题讨论】：

【解决方案1】：

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

在你的情况下，

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

应该工作

【讨论】：

当 x[i][j] 是标量时有效。现在我的 x[i][j] 是一个矩阵，这只是一个元素乘法，我想计算 x[i][0]= x[i][0].matmul(x[i][1 ]).matmul(x[i][2])……有什么办法可以实现吗？
在这种情况下，您可能必须使用内循环，但您可以通过使用 torch.bmm 进行乘法来避免外循环。 pytorch.org/docs/stable/torch.html#torch.bmm

【解决方案2】：

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

向量运算加速了这些计算，因为它们可以并行化。

【讨论】：