PyTorch中Affine grid和grid sample这两个兄弟函数的用法

前言

Affie Grid能用来实现ROI Pooling操作，在RCNN网络和STN（spatial transformation network）中常被用到。除了这些，他还有一个作用就是用来做仿射变换（affine transform）。
关于仿射变化，可以参考我的另一篇博客 仿射变换的基本原理
还有一点，affine grid基本上是和grid sample混用的。

最后需要说明的是，affine grid和grid sample这两个兄弟函数，很像opencv中的一对函数：getRotationMatrix2D 和 warpAffine。为啥这么说，看完本篇就知晓了。

正文

def affine_grid(theta, size, align_corners=None):
‘’’
theta: 一个 N*2*3的张量，N是batch size。
size： 是得到的网格的尺度，也就是希望仿射变换之后得到的图像大小

‘’’

回忆一下仿射变换，需要一个3*3的矩阵。而affine grid需要的theta是N*2*3的，其中的这个2*3就是仿射矩阵的前两行（因为第三行是涉及到透视变换的，和仿射变换无关，pytorch维护者就不管最后一行）。
再回忆一下，前两行都是控制什么的：

• 前两列涉及到旋转和缩放
• 最后一列涉及到平移
  这6个自由度，就组成了平面2d图像的仿射变换需要的变换矩阵。而affine grid，就是通过用户设置仿射变换的参数（theta），并且指定仿射变换之后得到的图像的size，得到了一个grid。这个grid输入到grid sample中，就完成了仿射变换。grid用来指定，新图像的每个位置，需要去源图像的哪个位置采样。

接下来做一下实验

    img = cv2.imread('./00000.jpg')
    img = cv2.resize(img, dsize=None, dst=None, fx=0.5, fy=0.5)
    size = img.shape[:2]
    size = (1, 3,) + size
    angle = 30 * np.pi / 180  # np默认角度为弧度制

    theta = np.array([np.cos(angle), np.sin(-angle), 0,
                      np.sin(angle), np.cos(angle), 0])  # 30度的旋转矩阵
    theta = theta.reshape(1, 2, 3)
    theta = torch.from_numpy(theta).float()  # 调整dtype
    grid = F.affine_grid(theta, size=torch.Size(size))  # 得到grid 用于grid sample
    # print(grid.min(), grid.max())
    # grid = torch.clamp(grid, -1, 1)
    img_tensor = torch.from_numpy(img).float().permute(2, 0, 1).unsqueeze(0)
    warp_img = F.grid_sample(img_tensor, grid).squeeze().permute(1, 2, 0).numpy()
    warp_img = np.clip(warp_img, 0, 255).astype(np.uint8)


    img_vis = np.concatenate((img, warp_img), axis=0) # 原图在上，affined在下
    cv2.imshow('res', img_vis)
    cv2.waitKey()

上面是原图，下面是affined图像。大概是被逆时针（正方向）旋转30度了。不过还是存在拉伸，不晓得为啥。

至于grid sample，请参见我的另一篇博客，理解他的原理
grid sample