SparK：首个在任何卷积神经网络上成功实现BERT/MAE风格预训练的模型  

这是ICLR论文《为卷积网络设计BERT：稀疏且分层的掩码建模》（arXiv:2301.03580）的官方实现，能够在任意CNN（如ResNet）上以BERT风格的自监督方式进行预训练。我们尽力使代码库简洁、易读、处于最先进水平，并且仅依赖最少的外部依赖。

https://user-images.githubusercontent.com/39692511/226858919-dd4ccf7e-a5ba-4a33-ab21-4785b8a7833c.mp4

🔥 最新消息

• 我们的ICLR海报页面上提供了简短的英文介绍！[📹录制视频、海报和幻灯片]。
• 5月11日将在OpenMMLab & ReadPaper（哔哩哔哩）再次进行直播！[📹录制视频]
• **4月27日（UTC+8 20:00）**将在OpenMMLab（哔哩哔哩）举行另一场直播！
• **3月22日（UTC+8 20:00）**将在极市平台（哔哩哔哩）举行另一场直播！[📹录制视频]
• 在TechBeat（将门创投）的分享也定于3月16日（UTC+8 20:00）！[📹录制视频]
• 我们很荣幸受Synced（微信公众号“机器之心机动组”）邀请，于**2月27日（UTC+0 11:00，UTC+8 19:00）**就SparK进行演讲，欢迎大家观看！[📹录制视频]
• 本工作已被ICLR 2023接收为Spotlight论文（前25%的优秀论文）。
• 其他报道：[Synced] [DeepAI] [TheGradient] [字节跳动] [CVers [量子位] [智源研究院] [机器之心机动组] [极市平台] [ReadPaper笔记]

🕹️ Colab可视化演示

请查看pretrain/viz_reconstruction.ipynb，以可视化SparK预训练模型的重建效果，例如：

我们还提供了pretrain/viz_spconv.ipynb，展示了密集卷积层中的“掩码模式消失”问题。

这里有哪些新亮点？

🔥 预训练的CNN性能超越预训练的Swin Transformer：

🔥 经过SparK预训练后，小型模型可胜过未预训练的大型模型：

🔥 所有模型均可受益，呈现规模效应：

🔥 生成式自监督预训练优于对比学习：

更多分析、讨论和评估，请参阅我们的论文。

待办事项列表

目录

• 预训练代码
• 自定义 CNN 模型的预训练教程（自定义 CNN 模型预训练教程）
• 自定义数据集的预训练教程（自定义数据集预训练教程）
• 预训练 Colab 可视化实验环境（重建，稀疏卷积）
• 微调代码
• Hugging Face 上的权重与可视化实验环境
• timm 中的权重

预训练权重（自监督；无解码器；可直接微调）

注意：网络定义方面，我们直接使用 timm.models.ResNet 和 官方 ConvNeXt。

reso.：图像分辨率；acc@1：ImageNet-1K 微调准确率（top-1）

架构	分辨率	acc@1	参数量	FLOPs	权重（自监督，不含 SparK 的解码器）
ResNet50	224	80.6	26M	4.1G	resnet50_1kpretrained_timm_style.pth
ResNet101	224	82.2	45M	7.9G	resnet101_1kpretrained_timm_style.pth
ResNet152	224	82.7	60M	11.6G	resnet152_1kpretrained_timm_style.pth
ResNet200	224	83.1	65M	15.1G	resnet200_1kpretrained_timm_style.pth
ConvNeXt-S	224	84.1	50M	8.7G	convnextS_1kpretrained_official_style.pth
ConvNeXt-B	224	84.8	89M	15.4G	convnextB_1kpretrained_official_style.pth
ConvNeXt-L	224	85.4	198M	34.4G	convnextL_1kpretrained_official_style.pth
ConvNeXt-L	384	86.0	198M	101.0G	convnextL_384_1kpretrained_official_style.pth

预训练权重（含 SparK 的 UNet 风格解码器；可用于重建图像）

架构	分辨率	acc@1	参数量	FLOPs	权重（自监督，含 SparK 的解码器）
ResNet50	224	80.6	26M	4.1G	res50_withdecoder_1kpretrained_spark_style.pth
ResNet101	224	82.2	45M	7.9G	res101_withdecoder_1kpretrained_spark_style.pth
ResNet152	224	82.7	60M	11.6G	res152_withdecoder_1kpretrained_spark_style.pth
ResNet200	224	83.1	65M	15.1G	res200_withdecoder_1kpretrained_spark_style.pth
ConvNeXt-S	224	84.1	50M	8.7G	cnxS224_withdecoder_1kpretrained_spark_style.pth
ConvNeXt-L	384	86.0	198M	101.0G	cnxL384_withdecoder_1kpretrained_spark_style.pth

安装与运行

我们强烈建议您使用 torch==1.10.0、torchvision==0.11.1 和 timm==0.5.4 进行复现。 请参阅 INSTALL.md 以安装所有 pip 依赖项。

• 三行代码加载预训练模型权重

# 首先下载我们的权重 `resnet50_1kpretrained_timm_style.pth`
import torch, timm
res50, state = timm.create_model('resnet50'), torch.load('resnet50_1kpretrained_timm_style.pth', 'cpu')
res50.load_state_dict(state.get('module', state), strict=False)     # 以防模型权重实际上保存在 state['module'] 中

• 预训练

  • 任何 ResNet 或 ConvNeXt 在 ImageNet-1K 上： 请参阅 pretrain/
  • 您自己的 CNN 模型： 请参阅 pretrain/，尤其是 pretrain/models/custom.py

• 微调

  • 任何 ResNet 或 ConvNeXt 在 ImageNet-1K 上： 请查看 downstream_imagenet/ 获取后续说明。
  • ResNet 在 COCO 上： 请参阅 downstream_d2/
  • ConvNeXt 在 COCO 上： 请参阅 downstream_mmdet/

致谢

我们参考了以下有用的代码库：

• BEiT、MAE、ConvNeXt
• timm、MoCoV2、Detectron2、MMDetection

许可证

本项目采用 MIT 许可证。更多详情请参阅 LICENSE。

引用

如果您觉得本项目有用，欢迎您给我们点个赞 ⭐，或在您的工作中引用我们 📖：

@Article{tian2023designing,
  author  = {Keyu Tian and Yi Jiang and Qishuai Diao and Chen Lin and Liwei Wang and Zehuan Yuan},
  title   = {为卷积网络设计 BERT：稀疏且分层的掩码建模},
  journal = {arXiv:2301.03580},
  year    = {2023},
}

SparK 快速上手指南

SparK 是首个成功将 BERT/MAE 风格的自监督预训练应用于任意卷积神经网络（CNN）（如 ResNet、ConvNeXt）的工具。通过稀疏且分层的掩码建模，它能让 CNN 在 ImageNet 等任务上获得超越 Swin-Transformer 的性能表现。

环境准备

为了复现最佳效果，官方强烈建议使用以下特定版本的依赖库：

• 操作系统: Linux (推荐)
• Python: 3.8+
• 核心依赖版本:
  • torch == 1.10.0
  • torchvision == 0.11.1
  • timm == 0.5.4

国内加速建议：在安装 Python 包时，推荐使用清华或阿里镜像源以提升下载速度。

安装步骤

1. 创建虚拟环境（可选但推荐）

python -m venv spark_env
source spark_env/bin/activate

2. 安装 PyTorch 及相关依赖

请根据您的 CUDA 版本选择合适的安装命令。以下示例使用清华镜像源安装指定版本的 PyTorch：

pip install torch==1.10.0 torchvision==0.11.1 -f https://download.pytorch.org/whl/torch_stable.html
# 或者使用国内镜像源（示例为清华源，请根据实际 cuda 版本调整）
# pip install torch==1.10.0+cu111 torchvision==0.11.1+cu111 -f https://mirror.tuna.tsinghua.edu.cn/pypi/web/simple

3. 安装其他依赖

克隆项目代码并安装其余依赖（主要是 timm）：

git clone https://github.com/keyu-tian/SparK.git
cd SparK
pip install timm==0.5.4 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install -r requirements.txt

基本使用

加载预训练模型

SparK 提供了多种架构（ResNet50/101/152/200, ConvNeXt-S/B/L）的自监督预训练权重。您可以直接在 3 行代码内加载这些权重进行微调或推理。

首先，从 Google Drive 或项目 README 提供的链接下载权重文件（例如 resnet50_1kpretrained_timm_style.pth）。

然后使用以下代码加载：

# 下载权重文件 resnet50_1kpretrained_timm_style.pth 到本地
import torch
import timm

# 创建模型并加载权重
res50 = timm.create_model('resnet50')
state = torch.load('resnet50_1kpretrained_timm_style.pth')
res50.load_state_dict(state['model']) # 注意：根据实际权重文件结构可能需要调整 key，通常直接 load_state_dict(state) 即可

# 设置为评估模式
res50.eval()

可视化重建效果

SparK 支持可视化模型对掩码图像的重建能力，以验证预训练效果。您可以直接运行项目提供的 Colab Notebook 或本地脚本：

# 查看重建可视化
jupyter notebook pretrain/viz_reconstruction.ipynb

# 查看稀疏卷积的掩码模式消失问题可视化
jupyter notebook pretrain/viz_spconv.ipynb

自定义模型预训练

若需在自己的数据集或自定义 CNN 架构上进行 SparK 预训练，请参考项目目录下的详细教程：

• 自定义模型: pretrain/#tutorial-for-pretraining-your-own-cnn-model
• 自定义数据集: pretrain/#tutorial-for-pretraining-your-own-dataset