如何有效地计算 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

当您使用两个循环来计算乘积时，复杂度为 n^2。想象一下，在反向传播过程中多次这样做，你的代码会变得很慢。向量运算加速了这些计算，因为它们可以并行化。