【pytorch】从yolo的make_grid理解torch.meshgrid、torch.stack

简述

yolo检测 make_grid创建网格代码如下，那么什么是torch.meshgrid?

def _make_grid(nx=20, ny=20):
        yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
        return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()

1、torch.meshgrid 创建行列坐标

torch.meshgrid 是 PyTorch 中的一个函数，这个函数通常用于创建坐标点的网格，以便进行一些网格上的操作，比如插值或者计算函数值。
例如，我们的目标是创建4行8列的坐标网格。

import torch
ny=4
nx=8
yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
print(yv)
print(xv)

'''
tensor([[0, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 1, 1, 1],
        [2, 2, 2, 2, 2, 2, 2, 2],
        [3, 3, 3, 3, 3, 3, 3, 3]])
tensor([[0, 1, 2, 3, 4, 5, 6, 7],
        [0, 1, 2, 3, 4, 5, 6, 7],
        [0, 1, 2, 3, 4, 5, 6, 7],
        [0, 1, 2, 3, 4, 5, 6, 7]])
'''

这个操作会形成以ny为行数，nx为列数的坐标网格，也就是4行8列。其中每行每列又按照torch.arange进行排序，也就是0-4行，0-8列。如下图所示。

可以看到yv实际上是一个纵坐标的网格，每个值都代表着自己所在的行，分别对应0-3行
xv代表着自己的列，也就是0-7列。在yolo检测的make_grid中，通常还有下个步骤，将两个值堆叠在一起。torch.stack。

2、torch.stack 结合行列坐标

torch.stack((xv, yv), 2)就是将上述行列坐标信息堆叠在第2个维度，这样子每个位置就有各自的行列坐标值。
(为了迎合yolo和w、h的顺序，x坐标堆叠在前面,y堆叠在后面)

tmp = torch.stack((xv, yv), 2)
print(tmp.shape)
print(tmp[1,2])
print(tmp[3,6])

'''
torch.Size([4, 8, 2])
tensor([2, 1])
tensor([6, 3])
'''

可以看到w=1,h=2对应的值就是[2,1]也就是第2行，第1列。
w=3,h=6对应的值就是[6,3]也就是第6行，第3列。每一个wh都对应各自的行列坐标。也就是通过grid和stack函数，每个行列都有属于自己的值了！

3、通过view函数扩展维度

out=tmp.view((1, 1, ny, nx, 2)).float()

也就是将3维向量扩展成5维，就是为了后续检测操作，没有什么需要特别说的。

至此，整个make_grid函数解释完毕。这个函数作用就是生成ny行，nx列的网格用于检测时候的xy坐标确定