简介

本仓库是用于视觉识别的上下文Transformer网络的官方实现。

CoT是一种统一的自注意力构建模块，可作为ConvNet中标准卷积的替代方案。因此，可以通过用CoT替换卷积来增强视觉骨干网络的上下文感知自注意力能力。

2021年3月25日至2021年6月5日：CVPR 2021 开放世界图像分类挑战赛

在CVPR 2021的开放世界图像分类挑战赛中获得第一名。（团队名称：VARMS）

使用方法

代码主要基于timm。

需求：

• PyTorch 1.8.0及以上版本
• Python 3.7
• CUDA 10.1及以上版本
• CuPy。

克隆仓库：

git clone https://github.com/JDAI-CV/CoTNet.git

训练

首先，下载ImageNet数据集。要在单节点8张GPU上对CoTNet-50进行350轮次的训练，运行以下命令：

python -m torch.distributed.launch --nproc_per_node=8 train.py --folder ./experiments/cot_experiments/CoTNet-50-350epoch

CoTNet（例如CoTNet-50）的训练脚本可以在cot_experiments文件夹中找到。

推理时间与准确率

无论是在默认还是高级训练设置下，CoTNet模型始终以更少的推理时间获得比其他视觉骨干网络更高的Top-1准确率。简而言之，CoTNet模型在推理时间和准确率之间实现了比现有视觉骨干网络更好的权衡。

ImageNet上的结果

名称	分辨率	参数量	FLOPs	Top-1准确率	Top-5准确率	模型
CoTNet-50	224	22.2M	3.3	81.3	95.6	GoogleDrive / Baidu
CoTNeXt-50	224	30.1M	4.3	82.1	95.9	GoogleDrive / Baidu
SE-CoTNetD-50	224	23.1M	4.1	81.6	95.8	GoogleDrive / Baidu
CoTNet-101	224	38.3M	6.1	82.8	96.2	GoogleDrive / Baidu
CoTNeXt-101	224	53.4M	8.2	83.2	96.4	GoogleDrive / Baidu
SE-CoTNetD-101	224	40.9M	8.5	83.2	96.5	GoogleDrive / Baidu
SE-CoTNetD-152	224	55.8M	17.0	84.0	97.0	GoogleDrive / Baidu
SE-CoTNetD-152	320	55.8M	26.5	84.6	97.1	GoogleDrive / Baidu

百度网盘的访问码为：cotn

CoTNet在下游任务中的应用

对于目标检测和实例分割任务，请参阅CoTNet用于目标检测和实例分割。

引用上下文Transformer网络

@article{cotnet,
  title={Contextual Transformer Networks for Visual Recognition},
  author={Li, Yehao and Yao, Ting and Pan, Yingwei and Mei, Tao},
  journal={arXiv preprint arXiv:2107.12292},
  year={2021}
}

致谢

感谢timm以及优秀的PyTorch团队所做的贡献。

CoTNet 快速上手指南

CoTNet（Contextual Transformer Networks）是一种统一的自注意力构建模块，可作为 ConvNet 中标准卷积的替代方案，旨在通过上下文自注意力增强视觉骨干网络。该模型在 ImageNet 分类任务中表现出色，并在推理时间与准确率之间取得了优异的平衡。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐)
• Python: 3.7+
• PyTorch: 1.8.0+
• CUDA: 10.1+
• 核心依赖: CuPy (用于加速计算)

提示：建议使用 Conda 创建虚拟环境以避免依赖冲突。国内用户可使用清华或中科大镜像源加速包安装。

安装步骤

1. 克隆仓库

首先从 GitHub 克隆官方代码库：

git clone https://github.com/JDAI-CV/CoTNet.git
cd CoTNet

2. 安装依赖

本项目主要基于 timm (PyTorch Image Models)。请安装必要的 Python 包：

pip install torch>=1.8.0 torchvision cupy-cuda10x timm

注意：cupy-cuda10x 需根据您的实际 CUDA 版本调整（如 cupy-cuda11x）。如果 pip 安装 CuPy 失败，请参考 CuPy 官方文档 进行编译安装。

基本使用

数据准备

下载并准备好 ImageNet 数据集。

训练模型

以下命令演示如何在单节点、8 张 GPU 上训练 CoTNet-50 模型（共 350 个 epoch）：

python -m torch.distributed.launch --nproc_per_node=8 train.py --folder ./experiments/cot_experiments/CoTNet-50-350epoch

• 训练脚本配置文件位于 cot_experiments 文件夹中，您可以根据需要修改超参数或切换其他变体（如 CoTNeXt-50, SE-CoTNetD-101 等）。

使用预训练模型

如果您仅需推理或使用预训练权重，可从下表获取模型文件（提供百度网盘和 Google Drive 下载）：

模型名称	分辨率	Top-1 Acc.	百度网盘链接 (提取码: cotn)
CoTNet-50	224	81.3%	下载
CoTNeXt-50	224	82.1%	下载
CoTNet-101	224	82.8%	下载
SE-CoTNetD-152	320	84.6%	下载

加载模型后，即可像使用标准 ResNet 一样将其集成到您的检测或分割任务中。针对目标检测和实例分割的专用代码，请访问 CoTNet 下游任务仓库。