Sampling point pairs from a grid in Pytorch
I need sample point pairs from a grid in PyTorch.

I have a tensor of size (1 x 500 x 1000). I also have a mask tensor of size(1 x 500 x 1000), denoting if a point is valid or not. I want to sample 200k point pairs from this grid. In other words, I want to get coordinates of sampled point pairs as a tensor of size (200k x 4), denoting (x1, y1, x2, y2) for all 200k point pairs. All points in the pairs should be valid points.

This will be repeated many times, so I need to have an efficient way of performing this procedure. What is an elegant way to implement this in PyTorch?

Not an expert here, but I did spend some time trying things out.
Turns out operating on 1D array is a lot faster (method two).

import time
import torch
class Timer():
    def __init__(self):
        pass
    def __enter__(self):
        self.time = time.time()
    def __exit__(self, *exc):
        print(f'time used: {time.time() - self.time:.2f}s')

# a = torch.rand([1,500,1000])
m = torch.randint(2, [1, 500, 1000]) # mask tensor
valid_len = (m==1).nonzero().size()[0] # number of valid points
rand_one = torch.randint(valid_len, [200000]) # sample 200k of random int
rand_two = torch.randint(valid_len, [200000]) # sample 200k of random int

# method one
m0 = m == 1 # mask of shape torch.Size([1, 500, 1000])
m0 = m0.nonzero() # valid points of shape torch.Size([valid_len, 3])
m0 = m0[:, 1:] # reshape to shape torch.Size([valid_len, 2])
with Timer():
    one0 = torch.index_select(m0, 0, rand_one) # take 200k valid points
    two0 = torch.index_select(m0, 0, rand_two) # take 200k valid points again
    coor0 = torch.cat([one0, two0], dim=1) # stack them up
# &gt;&gt;&gt; time used: 1.05s

# method two
m1 = m.reshape(-1) # reshape mask to torch.Size([500000])
m1 = m1==1 # mask of shape torch.Size([500000])
m1 = m1.nonzero() # valid points of shape torch.Size([valid_len, 1])
m1 = m1.reshape(-1) # valid points of shape torch.Size([valid_len])
with Timer():
    one1 = m1.take(rand_one) # take 200k valid points
    two1 = m1.take(rand_two) # again
    # transform them to coordinates and stack them up
    coor1 = torch.stack([one1 // 1000, one1 % 1000, two1 // 1000, two1 % 1000], dim=1)
# &gt;&gt;&gt; time used: 0.07s

assert torch.sum(coor0 == coor1) == 800000 # make sure consistent result 


cheers
