How to change some value of a tensor into zero according to another tensor's value in pytorch?

user_11955260 0
Viewed 762

I have two tensors: tensor a and tensor b. How can I change some value of tensor a according to the value of tensor b?

I know the codes following are right, but it runs pretty slow when the tensor is big. Is there any other method?

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

I guess torch.where would be faster I have measure in CPU here is result.

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)
user_60756990

For this exact use case also consider

a * (b <= 0.5)

which seems to be the fastest out of the following

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)
user_29802460
Related