通俗易懂的STN模块(二) 导读上一篇通俗易懂的STN模块

导读
上一篇通俗易懂的STN模块(一)中，我们详细介绍了STN中会使用到的几个模块，并且用pytorch和numpy来实现了，这篇文章我们将会利用pytorch来实现一个MNIST的手写数字识别并且将STN模块嵌入到CNN中
STN关键点解读
STN有一个最大的特点就是STN模块能够很容易地嵌入到CNN中，只需要进行非常小的修改即可。上一篇文章我们也说了STN拥有平移、旋转、剪切、缩放等不变性，而这一特点主要是依赖θ参数来实现的。刚开始的时候我还以为训练STN还需要准备θ标签数据，实际上并不需要。
【通俗易懂的STN模块(二)】当输入图片通过STN模块之后获得变换后的图片，然后我们再将变换后的图片输入到CNN网络中，通过损失函数计算loss ，然后计算梯度更新θ参数，最终STN模块会学习到如何矫正图片。
代码实现
导包

文章图片
导入相关的依赖包
定义网络结构

文章图片
嵌入STN模块的CNN
加载数据集

文章图片
训练模型

文章图片
评估模型

文章图片
将pytorch的tensor转换为numpy的array

文章图片
可视化STN变换图片

文章图片

文章图片
通过对比输入图片和经过STN变换后的图片能够很明显发现，经过STN之后能将旋转的图片进行明显的纠正。
参数设置

文章图片
主函数

文章图片
源码请参考我的博客「链接」
参考：https://pytorch.org/tutorials/intermediate/spatial_transformer_tutorial.html