利用MindSpore复现ICCV2021 Best Paper Swin Trasnformer

经过长达一个月的复现，终于成功利用MindSpore复现了SwinTransformer在imagenet上的分类精度，中间踩过很多的坑，这个帖子就作为复现SwinTransformer的记录贴，希望能对大家复现2021年这种充满训练Trick的论文有所帮助。

复现着复现着突然Swin就拿了最佳论文了，当时感觉也非常有意思，突然就在复现ICCV2021的最佳论文了，模型的效果的确很炸裂。

博客所有的相关代码已经上传到我的码云

修改完成后代码将会合入MindSpore的model的models主仓，有需要的同学可以自取。

数据篇——数据增强

由于SwinTransformer源码是基于PyTorch和timm完成的，其中的AutoAugment虽然是基本基于PIL库实现的，但是由于MindSpore本身图像库接口和timm和PyTorch存在一定的区别，非常不易自己实现。

因此笔者选择将timm的相关数据增强代码复制到MindSpore中并且基于numpy对其中PyTorch的一部分完整修改，同时也可以认为相当于是MindSpore的dataset扩充了一个可以即插即用的AutoAugment，相关的接口和PyTorch完全统一（反正也是cv来的），可以作为未来复现此类论文的一个基本模板，代码见我的仓库swin_transformer/data/data_utils文件夹

模型篇——混合精度

混合精度训练方法是通过混合使用单精度和半精度数据格式来加速深度神经网络训练的过程，同时保持了单精度训练所能达到的网络精度。混合精度训练能够加速计算过程，同时减少内存使用和存取，并使得在特定的硬件上可以训练更大的模型或batchsize。

这里我们主要针对MindSpore的混合精度和PyTorch（Apex）混合精度的区别做一些说明。

这里应该倒过来看，在笔者从模型训练、数据、性能角度都匹配了原基于PyTorch的SwinTransformer后，模型在10轮的时候会比同期PyTorch低1-2个点，这让人非常疑惑，因此最后联想到了混合精度。（不确定最后能不能靠O2跑到，一次三天实在等不起，就选择还是尽量同步源码了）

下面的理解是我基于MindSpore源码得到的，如果有错误希望指正。

首先是关于运算模式的一些定义：

• O0：纯FP32训练，可以作为accuracy的baseline；
• O1：混合精度训练（推荐使用），根据黑白名单自动决定使用FP16（GEMM, 卷积）还是FP32（Softmax）进行计算。
• O2：“几乎FP16”混合精度训练，不存在黑白名单，除了BatchNorm，几乎都是用FP16计算。
• O3：纯FP16训练，很不稳定，但是可以作为speed的baseline；

MindSpore拥有了O0、O2和O3，但是O1是缺省的，对于O2和O3，唯一的区别就是BN层是fp32还是fp16，那就是说如果不使用BN层的大多数ViT，O2和O3是一样的。

知道了这个区别之后，我们了解到Apex库中有对于fp16和fp32在O1模式下的名单，见
list
可以看到，实际在O1模式下，softmax、layernorm这类对精度影响较大的模块都会被保留在fp32，因此我们也选择将这些模块保存在fp32，进一步同步PyTorch。

关于转fp32和fp16的操作，选择直接去抄mindspore的amp.py文件：

print(f"=> using amp_level {
     args.amp_level}")
# 转换fp16
net.to_float(mstype.float16)
print(f"=> change {
     args.arch} to fp16")

# 转换fp32，反正swin就用了几个nn.Conv2d，主要还是nn.Dense，我就索性转fp32了
cell_types = (nn.GELU, nn.Softmax, nn.Conv2d, nn.Conv1d, nn.BatchNorm2d, nn.LayerNorm)
_do_keep_fp32(net, cell_types)
print(f"=> cast {
     cell_types} to fp32 back")
class OutputTo16(nn.Cell):
    "Wrap cell for amp. Cast network output back to float16"

    def __init__(self, op):
        super(OutputTo16, self)._