Как я могу эффективно вычислить тензор в 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.Представьте, что вы делаете это много раз во время обратного распространения, и ваш код замедляется.

Векторные операции ускоряют эти вычисления.

Answer 2

Это медленно, потому что вы используете два для циклов.

Вы можете использовать .prod См .: https://pytorch.org/docs/stable/torch.html#torch.prod

В вашем случае

x = torch.prod(x, dim=1) или x = x.prod(dim=1)

должно работать