2024-03-29 14:27:58 发布
网友
我有两个张量:张量a和张量b。如何根据张量b的值来改变张量a的值
我知道下面的代码是对的,但当张量很大时,它运行得相当慢。还有其他方法吗
import torch a = torch.rand(10).cuda() b = torch.rand(10).cuda() a[b > 0.5] = 0.
我猜torch.where会更快,我在CPU中进行了测量,结果如下
torch.where
import torch a = torch.rand(3**10) b = torch.rand(3**10)
%timeit a[b > 0.5] = 0. 852 µs ± 30.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit temp = torch.where(b > 0.5, torch.tensor(0.), a) 294 µs ± 4.51 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
此精确用例也考虑
a * (b <= 0.5)
以下哪一项似乎是最快的
In [1]: import torch ...: a = torch.rand(3**10) ...: b = torch.rand(3**10) In [2]: %timeit a[b > 0.5] = 0. 553 µs ± 17.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each) In [3]: a = torch.rand(3**10) In [4]: %timeit temp = torch.where(b > 0.5, torch.tensor(0.), a) ...: 49 µs ± 391 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each) In [5]: a = torch.rand(3**10) In [6]: %timeit temp = (a * (b <= 0.5)) 44 µs ± 381 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each) In [7]: %timeit a.masked_fill_(b > 0.5, 0.) 244 µs ± 3.48 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
我猜
torch.where
会更快,我在CPU中进行了测量,结果如下此精确用例也考虑
以下哪一项似乎是最快的
相关问题 更多 >
编程相关推荐