在pytorch中使用掩码进行赋值比for循环更快吗？

Question

我想问在pytorch中使用掩码分配是否比for循环更快。 例如，我们有以下两个代码块。

a = torch.tensor([-1,-1,1,1])
for i in range(len(a)):
    if a[i] > 0:
        a[i] = 0

a = torch.tensor([-1,-1,1,1])
a[a > 0] = 0

我很好奇哪个更快。有什么区别吗？

看起来我们需要两次遍历才能使用遮罩。但是在实现代码时确实有区别吗？

Answer 1

tensor

支持并行计算，因此它肯定比迭代快得多。您应该尽可能尝试使用

mask

。

在我的本地机器上快速测试一下，你可以自己尝试一下：

import torch
import time
import random

lst = [-1] * 5000 + [1] * 5000
random.shuffle(lst)

a = torch.tensor(lst)
b = torch.tensor(lst)

time_start = time.time()
for i in range(a.shape[0]):
    if a[i] > 0:
        a[i] = 0
print(f"Iteration takes: {time.time() - time_start}")

time_start = time.time()
b[b > 0] = 0
print(f"Mask takes: {time.time() - time_start}")

Iteration takes: 0.07204079627990723 \nMask takes: 0.011544942855834961