本次复现的论文是图像修复领域较火的Lama模型中的backbone,LaMa的backbone用在图像分类任务上也能有很好的效果。
原论文:https://arxiv.org/abs/2109.07161
源代码(用于图像修复的代码):https://github.com/saic-mdal/lama
复现的代码(用于图像分类代码):https://github.com/RooKichenn/General-lama-backbone(支持apex混合精度训练、mixup/cutmix、resume)
文章目录
一、LaMa整体结构
从图中可以看出,输入到网络的x'
是由真实图像x
和掩码m
在通道方向上进行拼接构成的,公式为x'=stack(x⊙m, m)
,拼接完成后,x.shape(3, 224, 224) -->x'.shape(4, 224, 224)
,进行下采样,在第一层下采样时用了一个7x7的快速傅里叶卷积(FFC
),LaMa的backbone整体使用的都是FFC
进行卷积,第一层使用大卷积核是为了获得全局的感受野,后两层下采样也是使用的FFC
,但是卷积核为3x3,经过下采样后,图像缩小了8倍,但是一般图像分类的模型都是下采样32倍,所以可以添加两个下采样模块,就是在原论文的基础上再下采样4倍,就变成了整体下采样32倍(我复现的是下采样8倍,32倍的还没试过)。下采样完成后,将特征图送入FFC ResNet Blocks
,在FFC ResNet Blocks
中特征图尺寸是不会变化的,主要目的就是为了让模型学习特征映射,后面的上采样我们不需要,我们只需有backbone,所以不做过多讲解,具体可以去阅读原论文。
二、 快速傅里叶卷积(FFC)
传统的卷积在网络的前几层并不能获得很大的感受野,这将浪费很多计算空间来对感受野进行建模,而且较小的感受野缺失全局的信息,于是作者使用了快速傅里叶卷积来获得全局的感受野。快速傅里叶卷积原论文地址:快速傅里叶卷积论文。FFC
用傅里叶变换到频域,在频域内做卷积,然后再逆变换回去。FFC
包含四个部分:局部-->局部
,局部-->全局
,全局-->局部
,全局-->全局
,前三部分都是使用标准3x3卷积实现的,只有最后一部分是使用频域变换。频域变换由两部分组成:Fourier Unit (FU)
和Local Fourier Unit (LFU)
。
Fourier Unit (FU)
伪代码如下图所示:
作者给出的伪代码介绍的非常清楚,其实就是先做傅里叶变换,在频域对实部和虚部做卷积,然后再逆变换回去。内部的卷积层是1×1卷积。
Local Fourier Unit (LFU)
目的是捕捉半全局信息,做法是将input feature map分为4个patch,分别使用FU即可。这样影响的是四分之一空间内的特征关联。就相当于是局部的nonlocal,想法合理。LFU 计算量比FU大,因为通道数多。实验说明在分类实验中,LFU还是很有用的。
三、FFC中local和global通道划分比例
FFC
原论文实验中给出了效果最好的比例,ratio=0.5,在LaMa模型中作者也使用的是ratio=0.5。
四、训练策略
在ImageNet上做预训练,也可以在自己的数据集上进行训练,只需要修改config.py中的类别个数即可。目前学习率还有各种超参数我都是默认使用的Swin Transformer
的超参,在ImageNet上做预训练开销太大,超参也没时间调。如果想训练自己的数据集,可以把学习率和优化器的参数适当调大一点。
代码开源地址:https://github.com/RooKichenn/General-lama-backbone