spatial transformer network (STN)

1.F.affine_grid
作用： 根据仿射变换矩阵，生成该变换相对应的逐像素偏移矩阵；
输入：theta：尺寸为(batch_size,2,3)的矩阵，表示每张图像的仿射变换矩阵
size:输出图像的尺寸
输出：根据theta变化得到的逐像素偏移矩阵

2.F.grid_sample 则是根据这个偏移矩阵将原始图片逐像素扭曲warp到目标位置处形成一个新的图片。
输入：img:要warp的图像
grid:逐像素偏移矩阵
输出：偏移后的图像

例子：

#将图像扭转45度
img = ToTensor()(img).unsqueeze(0)  # img→tensor 3维→4维
M = torch.tensor(                   # 仿射变换矩阵
    [[[0.7071, -0.7071, 0],
      [0.7071, 0.7071, 0]]]
)
grid = F.affine_grid(M, size=(1, 3, 500, 500))
warp_img = F.grid_sample(img, grid)             # 扭转图片
tensor_to_pil = ToPILImage()                    # tensor→img
tensor_to_pil(warp_img.data.squeeze(0)).show()  # 4维→3维

根据指定角度扭转图像
#feature.size()[0]就是batch size
            theta = torch.zeros((feature.size()[0], 2, 3), requires_grad=False, device=feature.device)
            with torch.no_grad():
                # 指定要旋转的角度
                angle = rotate_angle * math.pi / 180.0          
                #计算放射变换矩阵theta   
                theta[:, 0, 0] = torch.tensor(math.cos(angle), requires_grad=False, device=feature.device)
                theta[:, 0, 1] = torch.tensor(math.sin(-angle), requires_grad=False, device=feature.device)
                theta[:, 1, 0] = torch.tensor(math.sin(angle), requires_grad=False, device=feature.device)
                theta[:, 1, 1] = torch.tensor(math.cos(angle), requires_grad=False, device=feature.device)
			#根据放射变换矩阵得到逐像素偏移矩阵
            grid = F.affine_grid(theta, feature.size())
            #根据这个偏移矩阵将原始图片逐像素扭曲warp到目标位置处形成一个新的图片
            transformed_feature = F.grid_sample(feature, grid).to(feature.device)

STN可以参考blibli的李宏毅的教学