Conformer：局部特征与全局表征耦合的视觉识别网络

已被 ICCV21 接收！

本仓库基于 DeiT、timm 和 mmdetection 构建。

引言

在卷积神经网络（CNN）中，卷积操作擅长提取局部特征，但难以捕捉全局表征。而在视觉 Transformer 中，级联的自注意力模块能够捕获长距离的特征依赖关系，却往往导致局部特征细节的丢失。

本文提出了一种混合网络结构——Conformer，旨在结合卷积操作与自注意力机制的优势，以提升表征学习能力。Conformer 的核心是特征耦合单元（FCU），它以交互式的方式在不同分辨率下融合局部特征与全局表征。Conformer 采用并行结构，从而最大程度地保留了局部特征和全局表征。

实验表明，在参数量相当的情况下，Conformer 在 ImageNet 数据集上的性能比视觉 Transformer（DeiT-B）高出 2.3%。在 MSCOCO 数据集上，Conformer 分别在目标检测和实例分割任务中较 ResNet-101 提升了 3.7% 和 3.6% 的 mAP，展现了其作为通用骨干网络的巨大潜力。

Conformer 的基本架构如下所示：

我们还对比了 CNN（ResNet-101）、视觉 Transformer（DeiT-S）以及所提出的 Conformer 的特征图，如下所示。其中，Transformer 中的 patch embeddings 被重新调整为特征图以便可视化。CNN 会激活具有判别性的局部区域（例如，(a) 中的孔雀头部和 (e) 中的尾部），而 Conformer 的 CNN 分支则利用来自视觉 Transformer 的全局线索，从而激活完整的物体（例如，(b) 和 (f) 中的整只孔雀）。相比之下，视觉 Transformer 的局部特征细节有所损失（例如，(c) 和 (g)）。与此形成鲜明对比的是，Conformer 的 Transformer 分支在抑制背景的同时，保留了 CNN 中的局部特征细节（例如，(d) 和 (h) 中的孔雀轮廓比 (c) 和 (g) 更加完整）。

快速入门

安装

首先，安装 PyTorch 1.7.0+ 和 torchvision 0.8.1+，以及 pytorch-image-models 0.3.2：

conda install -c pytorch pytorch torchvision
pip install timm==0.3.2

数据准备

从 http://image-net.org/ 下载并解压 ImageNet 的训练和验证图像。目录结构应符合 torchvision datasets.ImageFolder 的标准布局，即训练数据位于 train/ 文件夹，验证数据位于 val/ 文件夹：

/path/to/imagenet/
  train/
    class1/
      img1.jpeg
    class2/
      img2.jpeg
  val/
    class1/
      img3.jpeg
    class/2
      img4.jpeg

训练与测试

训练

要在单节点 8 张 GPU 上对 Conformer-S 进行 300 个 epoch 的训练，运行以下命令：

export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
OUTPUT='./output/Conformer_small_patch16_batch_1024_lr1e-3_300epochs'

python -m torch.distributed.launch --master_port 50130 --nproc_per_node=8 --use_env main.py \
                                   --model Conformer_small_patch16 \
                                   --data-set IMNET \
                                   --batch-size 128 \
                                   --lr 0.001 \
                                   --num_workers 4 \
                                   --data-path /data/user/Dataset/ImageNet_ILSVRC2012/ \
                                   --output_dir ${OUTPUT} \
                                   --epochs 300

测试

要在单张 GPU 上对 Conformer-S 进行 ImageNet 测试，运行以下命令：

CUDA_VISIBLE_DEVICES=0, python main.py  --model Conformer_small_patch16 --eval --batch-size 64 \
                --input-size 224 \
                --data-set IMNET \
                --num_workers 4 \
                --data-path /data/user/Dataset/ImageNet_ILSVRC2012/ \
                --epochs 100 \
                --resume ../Conformer_small_patch16.pth

模型库

模型	参数量	MACs	Top-1 准确率	链接
Conformer-Ti	23.5 M	5.2 G	81.3 %	百度网盘(提取码: hzhm) 谷歌云盘
Conformer-S	37.7 M	10.6 G	83.4 %	百度网盘(提取码: qvu8) 谷歌云盘
Conformer-B	83.3 M	23.3 G	84.1 %	百度网盘(提取码: b4z9) 谷歌云盘

引用

@article{peng2021conformer,
      title={Conformer: Local Features Coupling Global Representations for Visual Recognition}, 
      author={Zhiliang Peng and Wei Huang and Shanzhi Gu and Lingxi Xie and Yaowei Wang and Jianbin Jiao and Qixiang Ye},
      journal={arXiv preprint arXiv:2105.03889},
      year={2021},
}

Conformer 快速上手指南

Conformer 是一种混合网络结构，旨在结合卷积神经网络（CNN）提取局部特征的优势与视觉 Transformer（ViT）捕捉全局依赖的能力。该模型在 ImageNet 分类、物体检测及实例分割任务中均表现出卓越性能。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐 Ubuntu)
• Python: 3.6+
• GPU: 支持 CUDA 的 NVIDIA 显卡
• 核心依赖:
  • PyTorch >= 1.7.0
  • torchvision >= 0.8.1
  • timm == 0.3.2

注意：本项目基于 DeiT, timm 和 mmdetection 构建，请确保安装指定版本的 timm 以避免兼容性问题。

安装步骤

1. 安装基础深度学习框架

推荐使用 Conda 进行环境管理。执行以下命令安装 PyTorch 和 torchvision：

conda install -c pytorch pytorch torchvision

2. 安装特定版本 timm

Conformer 依赖特定版本的 pytorch-image-models (timm)，请通过 pip 安装：

pip install timm==0.3.2

3. 获取代码

克隆本仓库到本地（假设您已获取源码仓库）：

git clone <repository_url>
cd <repository_directory>

基本使用

1. 数据准备

下载 ImageNet 数据集（训练集和验证集），并按照 torchvision.datasets.ImageFolder 的标准格式整理目录结构：

/path/to/imagenet/
  train/
    class1/
      img1.jpeg
    class2/
      img2.jpeg
  val/
    class1/
      img3.jpeg
    class2/
      img4.jpeg

2. 模型训练

以下是在单节点、8 张 GPU 上训练 Conformer-S 模型的示例命令（训练 300 epochs）：

export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
OUTPUT='./output/Conformer_small_patch16_batch_1024_lr1e-3_300epochs'

python -m torch.distributed.launch --master_port 50130 --nproc_per_node=8 --use_env main.py \
                                   --model Conformer_small_patch16 \
                                   --data-set IMNET \
                                   --batch-size 128 \
                                   --lr 0.001 \
                                   --num_workers 4 \
                                   --data-path /data/user/Dataset/ImageNet_ILSVRC2012/ \
                                   --output_dir ${OUTPUT} \
                                   --epochs 300

请根据实际路径修改 --data-path 参数。

3. 模型测试/评估

使用预训练权重在单张 GPU 上评估 Conformer-S 模型：

CUDA_VISIBLE_DEVICES=0, python main.py  --model Conformer_small_patch16 --eval --batch-size 64 \
                --input-size 224 \
                --data-set IMNET \
                --num_workers 4 \
                --data-path /data/user/Dataset/ImageNet_ILSVRC2012/ \
                --epochs 100 \
                --resume ../Conformer_small_patch16.pth

请将 --resume 后的路径替换为您下载的预训练模型文件路径。

4. 预训练模型下载

您可以从以下地址下载不同规模的预训练模型：

模型	参数量	Top-1 准确率	下载链接 (国内推荐)
Conformer-Ti	23.5 M	81.3%	百度网盘 (提取码: hzhm)
Conformer-S	37.7 M	83.4%	百度网盘 (提取码: qvu8)
Conformer-B	83.3 M	84.1%	百度网盘 (提取码: b4z9)