OpenCLIP

[论文] [引用] [Clip Colab] [Coca Colab]

欢迎使用 OpenAI CLIP（对比语言-图像预训练）的开源实现。

借助此代码库，我们已在多种数据源和计算预算上训练了多个模型，范围从小规模实验到包含在 LAION-400M、LAION-2B 和 DataComp-1B 等数据集上训练的大型运行。
我们训练的许多模型及其缩放特性在论文 reproducible scaling laws for contrastive language-image learning 中进行了详细研究。
我们训练的一些最佳模型及其零样本 ImageNet-1k 准确率如下所示，同时列出了 OpenAI 训练的 ViT-L 模型以及其他最先进的开源替代方案（均可通过 OpenCLIP 加载）。
关于我们完整预训练模型集合的更多细节，请参见此处，38 个数据集的零样本结果请参见此处。

包含更多模型特定细节的模型卡片可在 Hugging Face Hub 的 OpenCLIP 库标签下找到：https://huggingface.co/models?library=open_clip。

如果您觉得此仓库有用，请考虑引用。
如果您有任何其他请求或建议，欢迎提交 issue 或发送电子邮件。

请注意，src/open_clip/ 中的部分建模和分词器代码改编自 OpenAI 的官方仓库。

方法

图片来源：https://github.com/openai/CLIP

使用

pip install open_clip_torch

import torch
from PIL import Image
import open_clip

model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
model.eval()  # 默认模型处于训练模式，会影响一些具有 BatchNorm 或随机深度的模型
tokenizer = open_clip.get_tokenizer('ViT-B-32')

image = preprocess(Image.open("docs/CLIP.png")).unsqueeze(0)
text = tokenizer(["a diagram", "a dog", "a cat"])

with torch.no_grad(), torch.autocast("cuda"):
    image_features = model.encode_image(image)
    text_features = model.encode_text(text)
    image_features /= image_features.norm(dim=-1, keepdim=True)
    text_features /= text_features.norm(dim=-1, keepdim=True)

    text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1)

print("标签概率:", text_probs)  # 输出：[[1., 0., 0.]]

如果模型使用 timm 图像编码器（convnext、siglip、eva 等），请确保安装了最新版本的 timm。如果图像编码器出现“未知模型”错误，请升级 timm。

如果模型使用 transformers 分词器，请确保安装了 transformers。

另请参见此 [Clip Colab]。

要高效计算数十亿个嵌入，可以使用支持 openclip 的 clip-retrieval。

预训练模型

我们提供了一个简单的模型接口来实例化预训练和未训练的模型。
要查看哪些预训练模型可用，请使用以下代码片段。
关于预训练模型的更多细节，请参见此处。

>>> import open_clip
>>> open_clip.list_pretrained()

您可以在此表格中找到我们支持的模型的更多信息（例如参数数量、FLOPs）。

注意：许多现有的检查点使用原始 OpenAI 模型中的 QuickGELU 激活函数。在 PyTorch 的较新版本中，此激活函数实际上比原生的 torch.nn.GELU 效率低。模型默认现在使用 nn.GELU，因此对于 OpenCLIP 预训练权重，应使用带有 -quickgelu 后缀的模型定义。所有 OpenAI 预训练权重将始终默认使用 QuickGELU。也可以将非 -quickgelu 模型定义与使用 QuickGELU 的预训练权重一起使用，但会有准确率下降，对于较长的微调运行，这种下降可能会消失。
未来训练的模型将使用 nn.GELU。

加载模型

可以使用 open_clip.create_model_and_transforms 加载模型，如下例所示。模型名称和相应的 pretrained 键与 open_clip.list_pretrained() 的输出兼容。

pretrained 参数也接受本地路径，例如 /path/to/my/b32.pt。
您也可以通过这种方式从 huggingface 加载检查点。为此，请下载 open_clip_pytorch_model.bin 文件（例如，https://huggingface.co/laion/CLIP-ViT-L-14-DataComp.XL-s13B-b90K/tree/main），并使用 pretrained=/path/to/open_clip_pytorch_model.bin。

# pretrained 也接受本地路径
model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')

在分类任务上微调

此仓库专注于训练 CLIP 模型。要在下游分类任务（如 ImageNet）上微调已训练的零样本模型，请参见我们的另一个仓库：WiSE-FT。WiSE-FT 仓库包含我们关于零样本模型的鲁棒微调论文的代码，其中我们介绍了一种在微调零样本模型的同时保持分布偏移下鲁棒性的技术。

数据

要下载 webdataset 格式的数据集，我们推荐使用 img2dataset。

Conceptual Captions

参见 cc3m img2dataset 示例。

YFCC 和其他数据集

除了通过上述 CSV 文件指定训练数据外，我们的代码库还支持 webdataset，推荐用于更大规模的数据集。预期格式是一系列 .tar 文件。每个 .tar 文件应包含每个训练示例的两个文件，一个用于图像，一个用于相应的文本。两个文件应具有相同的名称但不同的扩展名。例如，shard_001.tar 可以包含诸如 abc.jpg 和 abc.txt 的文件。您可以在 https://github.com/webdataset/webdataset 了解更多关于 webdataset 的信息。我们使用每个包含 1,000 个数据点的 .tar 文件，这些文件是使用 tarp 创建的。

您可以从 Multimedia Commons 下载 YFCC 数据集。
与 OpenAI 类似，我们使用了 YFCC 的一个子集来达到上述准确率数字。
此子集中图像的索引位于 OpenAI 的 CLIP 仓库中。

训练 CLIP

安装

建议首先创建一个虚拟环境：

python3 -m venv .env
source .env/bin/activate
pip install -U pip

然后，您可以使用 pip install 'open_clip_torch[training]' 安装用于训练的 openclip。

开发

如果您想进行更改以贡献代码，可以克隆 openclip，然后在 openclip 文件夹中运行 make install（在创建虚拟环境之后）。

根据 https://pytorch.org/get-started/locally/ 安装 pip PyTorch。

您可以运行 make install-training 来安装训练依赖项。

测试

可以使用 make install-test 然后 make test 运行测试。

python -m pytest -x -s -v tests -k "training" 来运行特定测试。

针对特定 git 修订版或标签运行回归测试：
1. 生成测试数据
sh python tests/util_test.py --model RN50 RN101 --save_model_list models.txt --git_revision 9d31b2ec4df6d8228f370ff20c8267ec6ba39383
警告：这将调用 git 并修改您的工作树，但在数据生成后将重置为当前状态！\
在以这种方式生成测试数据时，请勿修改您的工作树。

2. 运行回归测试
   sh OPEN_CLIP_TEST_REG_MODELS=models.txt python -m pytest -x -s -v -m regression_test

单进程运行代码示例：

python -m open_clip_train.main \
    --save-frequency 1 \
    --zeroshot-frequency 1 \
    --report-to tensorboard \
    --train-data="/path/to/train_data.csv"  \
    --val-data="/path/to/validation_data.csv"  \
    --csv-img-key filepath \
    --csv-caption-key title \
    --imagenet-val=/path/to/imagenet/root/val/ \
    --warmup 10000 \
    --batch-size=128 \
    --lr=1e-3 \
    --wd=0.1 \
    --epochs=30 \
    --workers=8 \
    --model RN50

注意：imagenet-val 是 ImageNet 验证集的路径，用于零样本评估，而不是训练集！
如果您不想在整个训练过程中对 ImageNet 进行零样本评估，可以移除此参数。请注意，val 文件夹应包含子文件夹。如果不包含，请使用此脚本。

多 GPU 及更高规模

此代码已经过最多 1024 个 A100 的实战测试，并提供了多种分布式训练解决方案。我们包含了对 SLURM 集群的原生支持。

随着用于训练的设备数量增加，logit 矩阵的空间复杂度也会增加。使用朴素的 all-gather 方案，空间复杂度将为 O(n^2)。相反，如果使用标志 --gather-with-grad 和 --local-loss，复杂度可以有效地变为线性。这种修改会产生与朴素方法相同的一对一数值结果。

周期数

对于较大的数据集（例如 Laion2B），我们建议将 --train-num-samples 设置为低于完整周期数的值，例如 --train-num-samples 135646078 设置为周期的 1/16，并结合 --dataset-resampled 进行有放回采样。这样可以拥有频繁的检查点以更频繁地进行评估。

补丁丢弃

最近的研究表明，可以丢弃一半到四分之三的视觉标记，从而在不损失准确性的情况下将训练速度提高 2-3 倍。

您可以在视觉变换器配置中使用键 patch_dropout 设置此功能。

在论文中，他们最后在没有补丁丢弃的情况下进行了微调。您可以使用命令行参数 --force-patch-dropout 0. 来实现这一点。

多数据源

OpenCLIP 支持使用多个数据源，通过用 :: 分隔不同的数据路径。
例如，要在 CC12M 和 LAION 上训练，可以使用 --train-data "/data/cc12m/cc12m-train-{0000..2175}.tar::/data/LAION-400M/{00000..41455}.tar"。
在这些情况下，建议使用 --dataset-resampled。

默认情况下，模型从每个源看到样本的期望次数与该源的大小成正比。
例如，当在一个大小为 400M 和一个大小为 10M 的数据源上训练时，来自第一个源的样本在期望中被看到的可能性是 40 倍。

我们还支持对数据源进行不同的加权，通过使用 --train-data-upsampling-factors 标志。
例如，在上述场景中使用 --train-data-upsampling-factors=1::1 等同于不使用该标志，而 --train-data-upsampling-factors=1::2 等同于对第二个数据源进行两次上采样。
如果您想以相同的频率从数据源采样，上采样因子应与数据源的大小成反比。
例如，如果数据集 A 有 1000 个样本，数据集 B 有 100 个样本，您可以使用 --train-data-upsampling-factors=0.001::0.01（或类似地，--train-data-upsampling-factors=1::10）。

单节点

我们利用 torchrun 来启动分布式作业。以下命令在具有 4 个 GPU 的节点上启动一个作业：

cd open_clip/src
torchrun --nproc_per_node 4 -m open_clip_train.main \
    --train-data '/data/cc12m/cc12m-train-{0000..2175}.tar' \
    --train-num-samples 10968539 \
    --dataset-type webdataset \
    --batch-size 320 \
    --precision amp \
    --workers 4 \
    --imagenet-val /data/imagenet/validation/

多节点