Generating a meshgrid or itertools.product in pytorch

1

Hi all,

I’m trying to generate a meshgrid directly in pytorch…

The result I’m looking for, given two arrays of [0,1], is this:

[[0,0],
 [0,1],
 [1,0],
 [1,1]]

You can do this with np.meshgrid:

    grid = np.meshgrid(range(2), range(2), indexing='ij')
    grid = np.stack(grid, axis=-1)
    grid = grid.reshape(-1, 2)

and with itertools.product :

grid = list(itertools.product(range(2),range(2)))

But does anyone reckon how you’d do that directly in pytorch?

Thanks!

EDIT: Hm, this seems to work:

x = torch.Tensor([0,1])
torch.stack([x.repeat(2), x.repeat(2,1).t().contiguous().view(-1)],1)

x = torch.Tensor([0,1,2])
torch.stack([x.repeat(3), x.repeat(3,1).t().contiguous().view(-1)],1)

Full function for different sizes:

def generate_grid(h, w):
    x = torch.range(0, h-1)
    y = torch.range(0, w-1)
    grid = torch.stack([x.repeat(w), y.repeat(h,1).t().contiguous().view(-1)],1)
    return grid

grid = generate_grid(2,3)
# 0  0
# 1  0
# 0  1
# 1  1
# 0  2
# 1  2
#[torch.FloatTensor of size 6x2]
3 Likes