import torch
# Assuming that I have a tensor a.
a = torch.Tensor([-1,4,6,3,-1,-1])
# Then I need to select a random element(not -1) from a.
# How to do that efficiently?
This code should work:
a = torch.Tensor([-1,4,6,3,-1,-1])
valid_idx = (a!=-1).nonzero().view(-1)
choice = torch.multinomial(valid_idx.float(), 1)
a[valid_idx[choice]]
I’m not sure if it’s more efficient than just calling torch.randint(0, a.size(0), (1,))
in a while loop to skip the -1
entries. It might depend on the ratio between valid and invalid entries.
Hi, ptrblck
import torch
# Assuming that I have a tensor a.
a = torch.Tensor([[3,4, -1, -1,5,6],
[-1,5,-1,1,-1,-1],
[-1,4,4,2,3,0]])
# Then I need to select a random element(not -1) from a for each line
# How to do that efficiently?
# Eg.
# a_sample = [[4],[1],[3]]
In that case, this should work:
a = torch.Tensor([[3,4, -1, -1,5,6],
[-1,5,-1,1,-1,-1],
[-1,4,4,2,3,0]])
valid_idx = (a!=-1).nonzero()
choice = torch.multinomial(torch.arange(valid_idx.size(0)).float(), 1)
a[valid_idx[choice].squeeze().chunk(2)]
EDIT: Sorry, I didn’t realized you would like to sample from each line.
Will update the code in a minute.
EDIT2: I’m not sure, if we can avoid using a loop in this case:
valid_idx = (a!=-1).nonzero()
unique_rows = valid_idx[:, 0].unique()
valid_row_idx = [valid_idx[valid_idx[:, 0] == u] for u in unique_rows]
ret = []
for v in valid_row_idx:
choice = torch.multinomial(torch.arange(v.size(0)).float(), 1)
ret.append(a[v[choice].squeeze().chunk(2)])
ret = torch.stack(ret)
Hi, thank you for a great code snippet.
I found an error case where your code returns a Runtime Error.
When there is only one element in valid_idx
and that element is at the 0th position in a
it returns the following error:
RuntimeError: invalid argument 2: invalid multinomial distribution (sum of probabilities <= 0) at /Users/distiller/project/conda/conda-bld/pytorch_1556653464916/work/aten/src/TH/generic/THTensorRandom.cpp:343
My reproducible code is as follows:
a = torch.Tensor([-1,4,6,3])
indices = (a == -1).nonzero().view(-1)
choice = torch.multinomial(indices.float(), 1)
Your second suggestion works fine with this case.
I met a similar problem and just figured it out. You can achieve that with no loops:
a = torch.Tensor([[3,4, -1, -1,5,6],
[-1,5,-1,1,-1,-1],
[-1,4,4,2,3,0]])
b = torch.rand_like(a)
b[a == -1] = -1
a[torch.arange(a.shape[0]), b.argmax(dim=1)]
Edit: This requires generating many random numbers though, so it might be less efficient if the width of a
is very large.