Spatial Transformer Networks（空间转换器）及在MNIST中的应用

STN空间变换器在一些论文中会见到，而且因其简单有效、即插即用等特性，应用较多。为了充分理解论文和方便日后使用，这里记录一下STN以及应用在MNIST任务中。

1、Spatial Transformer Networks

空间转换器，空间变换器网络（简称STN）允许神经网络学习如何对输入图像进行空间变换，以增强模型的几何不变性。例如，它可以裁剪感兴趣的区域，缩放和校正图像的方向。它可能是一种有用的机制，因为 CNN 对于旋转和缩放以及更一般的仿射变换不是不变的。

在了解STN之前，你需要先学习仿射变换基础，可以看看图像仿射变换。
总之仿射变换就是原图*转换矩阵=仿射后的图，根据转换矩阵的不同，可以实现比如图像平移、缩放、旋转、翻转等等，而STN可以简单理解为通过CNN来自动学习转换矩阵，使得原图和转换矩阵运算后，能够被掰正。

比如在MNIST分类中，插入STN模块的效果，STN首先会将图像掰正，然后再次进行分类。

2、结构

上图就是STN模块，首先从UV的右侧面看出，图像经过该模块被纠正（旋转）了。

其代码如下：

def stn(self, x):
    """
    该部分经过卷积和全连接层，从原图拟合出用于仿射变换的转换矩阵，其shape=(2,3)
    :param x: 原图，shape=(1,28,28)
    :return:  仿射变换（掰正）后的图，shape=(1,28,28)
    """
    xs = self.localization(x)
    xs = xs.view(-1, 10 * 3 * 3)
    theta = self.fc_loc(xs)
    theta = theta.view(-1, 2, 3)    # 转换矩阵

    grid = F.affine_grid(theta, x.size())
    x = F.grid_sample(x, grid)

    return x

2.1 Localization net

注意，在代码中Localization net是包括self.localization和self.fc_loc两部分。

该部分完成从输入U中提取特征，拟合出变换矩阵参数θ，具体结构为：

2.2 Grid generator

其实这部分就只有一行代码：

grid = F.affine_grid(theta, x.size())

通过theta和希望变换后的尺寸s.size()来产生grid。关于这部分其实就是一个API的使用，可以百度一下就可。

2.3 Sampler

这部分也只有一行代码：

x = F.grid_sample(x, grid)

通过的得到的grid来对原图x进行变换，最终得到变换后的图x。

2.2 和2.3 这一步其实opencv也提供了warpAffine这个API来进行仿射变换，但是因为模型训练时tensor最好在GPU上运算，所以直接使用pytorch提供的affine_grid和grid_sample来进行仿射变换。这两个已经不属于深度学习的范畴了，就是稍微复杂一点的普通运算，百度一下即可。

3、实验

实验部分很简单，就是MNIST数字分类。只不过多加了STN模块。

实验 整体结构如下图所示，上面的STN负责将U变换掰正为V，下面的MNIST分类就是简单的CNN网络，这在很多关于MNIST入门教程中都能见到。

或者说就是在普通MNIST分类的CNN网络中，插入了STN模块，实现在使用CNN分类前先矫正MNIST图像数字的功能。

3.1 STN

STN部分将输入的图片U掰正并输出图片V，图片U和图片V的尺寸是一样的，都是$1\ast 28\ast 28$，下图展示实验训练100epochs后，U（Dataset Images）和V（Transformed Images），可以看到此时STN已经起作用了。图中一些歪七倒八的数字已经被”掰正“了。

3.2 MNIST分类

底下这部分就是卷积层+全连接层的简单分类网络，这不用多说了，MNIST就是深度学习的“Hello World！”。

3.3 效果展示

Spatial Transformer Networks对MNIST的“纠正”效果（epoch=20时的效果）

Test Acc曲线（增加epoch=100，能达到99%+）

4、代码

整个项目放在thgpddl/SpatialTransformerNetworks这里。

5、思考

在本实验中，MNIST中的一些样本通过STN被“掰正”，但是我们知道仿射变换还可以实现比如旋转、平移、缩放、裁剪等效果，那么是否在某些任务中，STN的作用可能时其他仿射效果呢？比如某任务中一些样本图像中目标很小，是否插入STN后的效果是将图像缩放至目标很大呢，这样不就相当于“注意力”了吗？

比如在论文阅读：Spatial Transformer Networks中，就有两个并行的STN，从效果来看，确实有放大的作用（从视觉上来看也可以是平移裁剪等实现），不就相当于一种注意力码?

6、引用

SPATIAL TRANSFORMER NETWORKS TUTORIAL