零样本模型的鲁棒微调

本仓库包含论文《零样本模型的鲁棒微调》（Robust fine-tuning of zero-shot models）的代码，该论文由 Mitchell Wortsman*、Gabriel Ilharco*、Jong Wook Kim、Mike Li、Simon Kornblith、Rebecca Roelofs、Raphael Gontijo-Lopes、Hannaneh Hajishirzi、Ali Farhadi、Hongseok Namkoong 和 Ludwig Schmidt 共同撰写，发表于 arXiv：https://arxiv.org/abs/2109.01903。

简而言之：我们在不增加微调或推理阶段计算开销的情况下，对零样本模型进行微调，同时保持或提升其 OOD 准确率。

摘要

大型预训练模型，如 CLIP 或 ALIGN，在执行零样本推理时（即无需针对特定数据集进行微调），能够在多种数据分布上提供一致的准确率。尽管现有的微调方法可以显著提升分布内准确率，但往往会降低分布外的鲁棒性。为解决这一矛盾，我们提出了一种简单而有效的方法来提升鲁棒性：将零样本模型和微调后模型的权重进行集成（WiSE-FT）。与标准微调相比，WiSE-FT 在保持高分布内准确率的同时，大幅提升了分布外准确率。在 ImageNet（分布内）以及五个衍生的分布偏移数据集上，WiSE-FT 相较于先前工作，将分布外准确率提高了 4 至 6 个百分点（pp），同时使分布内准确率提升了 1.6 个百分点。此外，WiSE-FT 在另外六组多样化的分布偏移数据集上也实现了类似的鲁棒性提升（2 至 23 个百分点），并且在七个常用的迁移学习数据集上，相比标准微调，分布内准确率提升了 0.8 至 3.3 个百分点。这些改进均未带来额外的计算成本。

总结图

代码

概述

WiSE-FT 可以在标准微调的基础上，通过几行代码实现，如下所示。更多细节请参阅 src/wise_ft.py。

# 加载模型
zeroshot = ImageClassifier.load(zeroshot_checkpoint)
finetuned = ImageClassifier.load(finetuned_checkpoint)
theta_0 = zeroshot.state_dict()
theta_1 = finetuned.state_dict()

# 确保检查点兼容
assert set(theta_0.keys()) == set(theta_1.keys())

# 使用混合系数 alpha 在检查点之间插值
theta = {
    key: (1-alpha) * theta_0[key] + alpha * theta_1[key]
    for key in theta_0.keys()
}

# 根据新权重更新模型
finetuned.load_state_dict(theta)

# 评估
evaluate(finetuned, args)

安装依赖

conda env create
conda activate wiseft

将目录添加到 PYTHONPATH：

cd wise-ft
export PYTHONPATH="$PYTHONPATH:$PWD"

下载数据

必要时，请参考 datasets.md，了解如何下载数据集。

运行 WiSE-FT

当零样本模型和微调后模型均已准备就绪时，示例命令如下：

python src/wise_ft.py   \
    --eval-datasets=ImageNet,ImageNetV2,ImageNetR,ImageNetA,ImageNetSketch  \
    --load=models/zeroshot.pt,models/finetuned.pt  \
    --results-db=results.jsonl  \
    --save=models/wiseft  \
    --data-location=~/data \
    --alpha 0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0

从头开始使用 ViT-B/32 运行 WiSE-FT 的示例命令如下：

python src/wise_ft.py   \
    --train-dataset=ImageNet  \
    --epochs=10  \
    --lr=0.00003  \
    --batch-size=512  \
    --cache-dir=cache  \
    --model=ViT-B/32  \
    --eval-datasets=ImageNet,ImageNetV2,ImageNetR,ImageNetA,ImageNetSketch  \
    --template=openai_imagenet_template  \
    --results-db=results.jsonl  \
    --save=models/wiseft/ViTB32  \
    --data-location=~/data \
    --alpha 0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0

注意：标志 --freeze-encoder 用于控制是仅微调线性分类器，还是对所有权重进行端到端微调。

结果绘图

生成散点图的示例命令如下：

python src/scatter_plot.py  \
    --eval-datasets=ImageNetV2,ImageNetR,ImageNetA,ImageNetSketch  \
    --results-db=results.jsonl  \
    --save plots

以下展示了使用 ViT-B/16 运行上述命令时的预期效果（模型可在此下载：链接）：

       

       

引用

如果您觉得本仓库有所帮助，请考虑引用以下文献：

@article{wortsman2021robust,
  title={Robust fine-tuning of zero-shot models},
  author={Wortsman, Mitchell and Ilharco, Gabriel and Kim, Jong Wook and Li, Mike and Kornblith, Simon and Roelofs, Rebecca and Gontijo-Lopes, Raphael and Hajishirzi, Hannaneh and Farhadi, Ali and Namkoong, Hongseok and Schmidt, Ludwig},
  journal={arXiv preprint arXiv:2109.01903},
  note={\url{https://arxiv.org/abs/2109.01903}},
  year={2021}
}

wise-ft 快速上手指南

wise-ft (WiSE-FT) 是一种用于零样本模型（如 CLIP）的鲁棒微调方法。它通过将零样本模型权重与微调后的模型权重进行混合（Ensembling），在无需额外计算成本的情况下，显著提升模型的分布外（OOD）鲁棒性，同时保持甚至提高分布内（ID）的准确率。

环境准备

• 系统要求: Linux 或 macOS 环境（Windows 需使用 WSL）。
• 前置依赖:
  • Python 3.8+
  • Conda (推荐用于环境管理)
  • GPU (推荐 NVIDIA GPU 以加速训练和推理)
• 网络建议: 由于需要下载预训练模型和数据集，建议确保网络连接通畅。如有条件，可配置国内镜像源加速 pip 或 conda 包的安装。

安装步骤

1. 克隆仓库并创建环境 首先克隆代码库，并使用项目提供的配置文件创建 Conda 环境：

   git clone https://github.com/mlfoundations/wise-ft.git
   cd wise-ft
   conda env create
   conda activate wiseft
   

2. 配置环境变量 将当前目录添加到 PYTHONPATH，以便 Python 能正确导入模块：

   export PYTHONPATH="$PYTHONPATH:$PWD"
   

3. 准备数据集 根据任务需求下载数据集（如 ImageNet 及其变体）。具体下载说明请参考项目内的 datasets.md 文件。

   提示：国内用户可尝试通过阿里云盘、百度网盘等渠道寻找公开数据集资源，或使用学术加速服务下载。

基本使用

WiSE-FT 的核心逻辑非常简单：加载零样本模型和微调模型的权重，按系数 $\alpha$ 进行线性插值混合。

场景一：已有模型权重，直接混合评估

如果你已经拥有了零样本模型（zeroshot）和微调后模型（finetuned）的检查点文件，可以直接运行以下命令进行混合评估。该命令会遍历不同的 $\alpha$ 值（从 0.0 到 1.0）以寻找最佳效果：

python src/wise_ft.py   \
    --eval-datasets=ImageNet,ImageNetV2,ImageNetR,ImageNetA,ImageNetSketch  \
    --load=models/zeroshot.pt,models/finetuned.pt  \
    --results-db=results.jsonl  \
    --save=models/wiseft  \
    --data-location=~/data \
    --alpha 0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0

场景二：从零开始训练并应用 WiSE-FT

如果你需要从头开始微调模型（例如使用 ViT-B/32），可以使用以下命令。该流程会自动执行微调，并在不同 $\alpha$ 值下评估混合后的模型效果：

python src/wise_ft.py   \
    --train-dataset=ImageNet  \
    --epochs=10  \
    --lr=0.00003  \
    --batch-size=512  \
    --cache-dir=cache  \
    --model=ViT-B/32  \
    --eval-datasets=ImageNet,ImageNetV2,ImageNetR,ImageNetA,ImageNetSketch  \
    --template=openai_imagenet_template  \
    --results-db=results.jsonl  \
    --save=models/wiseft/ViTB32  \
    --data-location=~/data \
    --alpha 0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0

参数说明：

• --freeze-encoder: 可选标志。若添加此标志，则仅微调线性分类器；若不添加，则对整个模型进行端到端微调。
• --alpha: 混合系数。0 代表纯零样本模型，1 代表纯微调模型，中间值代表两者的权重混合。

可视化结果

实验完成后，可使用以下命令生成散点图以直观展示分布内与分布外准确率的变化关系：

python src/scatter_plot.py  \
    --eval-datasets=ImageNetV2,ImageNetR,ImageNetA,ImageNetSketch  \
    --results-db=results.jsonl  \
    --save plots