CVPR 2021 | 超越卷积，自注意力模型HaloNet准确率实现SOTA

冷不防 2022-08-30 04:20 134阅读 0赞

##  ##

[https://github.com/lucidrains/halonet-pytorch][https_github.com_lucidrains_halonet-pytorch]

import torch
    from halonet_pytorch import HaloAttention
    
    if __name__ == '__main__':
        model = HaloAttention(dim=256,  # dimension of feature map
            block_size=8,  # neighborhood block size (feature map must be divisible by this)
            halo_size=4,  # halo size (block receptive field)
            dim_head=64,  # dimension of each head
            heads=4  # number of attention heads
        )#.cuda()
    
        input = torch.randn(1, 256, 32, 32)#.cuda()
        model(input)  # (1, 512, 32, 32)
    
        model_path = "dicenet.pth"
        torch.save(model.state_dict(), model_path)
    
        import os
    
        fsize = os.path.getsize(model_path)
        fsize = fsize / float(1024 * 1024)
    
        print(f"model size {round(fsize, 2)} m")
    
        for i in range(15):
            t1 = time.time()
            loc = model(input)
    
            cnt = time.time() - t1
            print(cnt, loc.size())

尽管该研究的模型使用自注意力代替卷积来捕获像素之间的空间交互，但它们采用了现代卷积神经网络（CNN）的一些重要架构特征。并且，和卷积神经网络一样，该研究计算多尺度特征层次结构，这些层次结构可以在局部化和实例分割等任务上以多种尺寸进行目标检测。

因此，该研究开发了一个 strided 自注意力层，它是 strided 卷积的自然扩展。为了在无法引起全局注意力的较大分辨率下处理计算成本，研究者遵循局部处理（也是卷积和自然感知系统的核心）的通用原则，并使用自注意力的空间受限形式。下图 2 为注意力下采样层的工作流程：

该方法没有使用平移等价性（translational equivariance）来代替更好的硬件利用率，从而改善了速度和准确率之间的权衡。尽管使用的是局部注意力，但每个像素的感受野却非常大（达到了 18×18），并且更大的感受野有助于处理更大的图像。

此外，研究者还介绍了用于视觉任务的自注意力，并描述了如何放松平移等价性，以有效地将局部自注意力映射到硬件。

实验结果

每个 HaloNet 模型（H0–H7）都是通过连续提升表 2 中定义的超参数的值来设计的。研究者后续还会进行比肩 EfficientNet 的更大 HaloNet 模型的训练和评估工作。

可与 SOTA 卷积模型媲美

该研究在 ImageNet 基准上训练了 HaloNet 模型，批大小是 4096，学习率是 1.6，线性预热了 10 个 epoch，然后进行了余弦衰减。模型用 Nesterov 的加速梯度（Accelerated Gradient）训练 350 个 epoch，并使用 dropout、权重衰减、RandAugment 和随机深度进行了正则化。

研究者验证了自注意力与视觉卷积各自的优势，并进一步理解了自注意力视觉架构的最佳设计方式。

实验结果如上表所示，带有正则化的 HaloNet 比 ResNet 的性能增益更多，但是对 ResNet 中重要的架构化模块却没有显著改进。

HaloNet 架构

研究者探究了放松平移等价性的影响以及邻域窗口与光晕大小的关系。

如下图 5 所示，放松平移等价性能够提升准确率。

从下图 6 可以发现随着窗口大小的增加，准确率会不断提高。特别是，将窗口大小从 6×6 增大到 12×12，准确率会提升 1.3%。这些结果表明，增加窗口大小可以用于扩展模型，而无需增加参数数量，这可能对生产环境有利。

卷积 - 注意力混合改善了速度 - 准确率权衡

在一组控制变量实验中，该研究用卷积替代了自注意力，以了解当前注意力层最有利的地方。

在性能最优的模型（HaloNet H7）的每个阶段中，表 4 展示了借助 SE 模块用卷积替代注意力层的结果。除最后一个阶段外，其他所有阶段都进行卷积会产生最快的模型，尽管 top-1 准确率会显著降低（1％）。在卷积和注意力上分别进行分配可最小程度地降低预测准确率，同时显著改善训练和推断过程。未来研究者还将对改进的混合模型进行进一步的研究。

[https_github.com_lucidrains_halonet-pytorch]: https://github.com/lucidrains/halonet-pytorch