torch.meshgrid() 函数理解

总是能遇到这个函数，又总是能忘记它原理。
其本质就是不断地repeat来达到想要的尺寸。
官方介绍

grid_x = torch.range(3)
grid_y = torch.range(4)
grid_z = torhc.range(5)

grid_x, grid_y, grid_y = torch.meshgrid([grid_x, grid_y, grid_z])
1
2
3
4
5

输出：grid_x, grid_y, grid_z尺寸均为$3\ast 4\ast 5$。

• grid_x[:,0,0] = [1, 2, 3]
• grid_y[0,:,0] = [1, 2, 3, 4]
• grid_x[0,0,:] = [1, 2, 3, 4, 5]

处理：.reshape([-1, 1])时，沿着z, y, x展开。

• grid_x = grid_x.reshape([-1, 1])：20个0，20个1，20个2
• grid_y = grid_y.reshape([-1, 1])：[5个0，5个1，5个2，5个3] * 3
• grid_z = grid_z.reshape([-1, 1])：[[0, 1, 2, 3, 4] * 4] *5

shape is [60, 1]

最终：
torch.cat([grid_x, grid_y, grid_z], dim=1)：

…

shape is [60, 3]