时间戳

首页 > 解决方案 > 如何根据pytorch中另一个张量的值将某个张量的值更改为零？

python - 如何根据pytorch中另一个张量的值将某个张量的值更改为零？

问题描述

我有两个张量：张量 a 和张量 b。如何根据张量 b 的值更改张量 a 的某些值？

我知道下面的代码是正确的，但是当张量很大时它运行得非常慢。还有其他方法吗？

import torch
a = torch.rand(10).cuda()
b = torch.rand(10).cuda()
a[b > 0.5] = 0.

标签： pythonpytorch

解决方案

我想torch.where我在 CPU 中测量的结果会更快。

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)

推荐阅读