【numpy&pytorch】meshgrid的shape

简单粗暴，直接上代码。

numpy

import numpy as np


xx = np.arange(3, 7)  # 3-6
yy = np.arange(0, 10) # 0-9

grid_x, grid_y = np.meshgrid(xx, yy)

print(grid_x, grid_x.shape)  # [10, 4], e.g. [n_y, n_x]
print(grid_y, grid_y.shape)  # [10, 4]

[[3 4 5 6]
 [3 4 5 6]
 [3 4 5 6]
 [3 4 5 6]
 [3 4 5 6]
 [3 4 5 6]
 [3 4 5 6]
 [3 4 5 6]
 [3 4 5 6]
 [3 4 5 6]] (10, 4)
[[0 0 0 0]
 [1 1 1 1]
 [2 2 2 2]
 [3 3 3 3]
 [4 4 4 4]
 [5 5 5 5]
 [6 6 6 6]
 [7 7 7 7]
 [8 8 8 8]
 [9 9 9 9]] (10, 4)

torch

imoprt torch

xx = torch.arange(3, 7)  # 3-6
yy = torch.arange(0, 10) # 0-9

grid_x, grid_y = torch.meshgrid(xx, yy)

print(grid_x, grid_x.shape)  # [4, 10], e.g. [n_x, n_y]
print(grid_y, grid_y.shape)  # [4, 10]

tensor([[3, 3, 3, 3, 3, 3, 3, 3, 3, 3],
        [4, 4, 4, 4, 4, 4, 4, 4, 4, 4],
        [5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
        [6, 6, 6, 6, 6, 6, 6, 6, 6, 6]]) torch.Size([4, 10])
tensor([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]]) torch.Size([4, 10])

结论

• numpy, meshgrid(xx, yy) ⇒ shape (n_y, n_x)
• torch, meshgrid(xx, yy) ⇒ shape (n_x, n_y)

千万注意！为了不搞混，建议记住一种，另一种通过tensor-numpy转换来得到。其实numpy.meshgrid才是正常人的理解，因为 $ny$ 是行， $nx$ 是列。使用torch.meshgrid时注意关键字位置应该为grid_y, grid_x= torch.meshgrid(yy, xx)才能得到想要的 $ny\times nx$ 网格。
因此时常会见到以下操作：

import torch

grid_y, grid_x = torch.meshgrid([torch.arange(h), torch.arange(w)])  # shape: (h, w)
grid_xy = torch.stack([grid_x, grid_y], dim=-1).float()  # shape: (h, w, 2), 2=(x, y)