PyTorch meshgrid 生成网格坐标

torch.meshgrid(*tensors, indexing=None)

使用输入的 1-D 张量创建网格坐标

示例

1. indexing=‘xy’

import torch

a = torch.arange(3)
b = torch.arange(3, 6)

x, y = torch.meshgrid(a, b, indexing='xy')
print('x =', x, sep='\n')
print('y =', y, sep='\n')
1
2
3
4
5
6
7
8

输出:

x =
tensor([[0, 1, 2],
        [0, 1, 2],
        [0, 1, 2]])
y =
tensor([[3, 3, 3],
        [4, 4, 4],
        [5, 5, 5]])
1
2
3
4
5
6
7
8

输入的 a, b 张量分别作为横坐标和纵坐标, 组合形成 3x3 的坐标元组, 每个坐标元组的第一个元素是横坐标, 第二个元素是纵坐标

按位置提取所有坐标元组的第一个元素, 得到输出 x, 提取第二个元素得到输出 y

torch.meshgrid(*tensors) currently has the same behavior as calling numpy.meshgrid(*arrays, indexing=’ij’).
In the future torch.meshgrid will transition to indexing=’xy’ as the default.

在未来的 PyTorch 版本, indexing 的默认参数将改为 ‘xy’, 等价于 Numpy.meshgrid(*arrays) 函数

2. indexing=‘ij’

现版本的 indexing 参数默认为 ‘ij’, 等价于 Numpy.meshgrid(*arrays, indexing=‘ij’) 函数

>>> import torch
>>> a = torch.arange(3)
>>> b = torch.arange(3, 6)
>>> x, y = torch.meshgrid(a, b, indexing='ij')
>>> x
tensor([[0, 0, 0],
        [1, 1, 1],
        [2, 2, 2]])
>>> y
tensor([[3, 4, 5],
        [3, 4, 5],
        [3, 4, 5]])
1
2
3
4
5
6
7
8
9
10
11
12

与 indexing=‘xy’ 不同, indexing=‘ij’ 时, 输入的 a 张量 (0 1 2) 是作为纵坐标的, 而且坐标元组中纵坐标排在第一个元素, 输出的 x 张量对应纵坐标元素

从输出结果来看, 两种 indexing 参数生成的张量互为转置关系

其他

以上示例创建的都是二维网格坐标, 如果输入 3 个 1-D 张量, 得到是三维网格坐标