OpenCLIP 开源实现与训练实战指南

然后安装核心库：

pip install open_clip_torch

如果需要进行训练，还需安装训练依赖：

pip install 'open_clip_torch[training]'

基础使用示例

加载预训练模型并进行图像文本特征提取是核心流程。这里有一个简单的推理示例：

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()  # 默认是训练模式，评估时需设为 eval

# 获取分词器
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.cuda.amp.autocast():
    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("Label probs:", text_probs)  # 输出如 [[1., 0., 0.]]

注意：为了高效计算数十亿个嵌入，可以使用支持 openclip 的专用工具。

模型管理

预训练模型列表

可以通过以下代码查看可用的预训练模型：

import open_clip
open_clip.list_pretrained()

关于模型参数量、FLOPs 等详细信息，建议查阅官方文档。

重要提示：许多现有检查点使用原始 OpenAI 模型中的 QuickGELU 激活函数。这在 PyTorch 新版本中效率较低。现在模型默认值为 nn.GELU，因此对于 OpenCLIP 预训练权重，应使用带 -quickgelu 后缀的模型定义。所有 OpenAI 预训练权重始终默认为 QuickGELU。虽然也可以使用非 -quickgelu 模型定义配合 QuickGELU 权重，但精度可能会下降，尤其是在微调时。

加载模型

使用 open_clip.create_model_and_transforms 加载模型，名称和 pretrained 键需与 list_pretrained() 的输出兼容。

# 支持本地路径或 HuggingFace 检查点
model, _, preprocess = open_clip.create_model_and_transforms(
    'ViT-B-32', 
    pretrained='/path/to/my/b32.pt'  # 或 laion2b_s34b_b79k
)

若从 HuggingFace 加载，请下载 open_clip_pytorch_model.bin 文件并使用 pretrained=/path/to/open_clip_pytorch_model.bin。

训练指南

单进程训练

这是最基础的启动方式，适合调试或小规模实验：

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 指向的是用于零样本评估的验证集路径，而非训练集。如果不需要在整个训练过程中进行零样本评估，可以移除该参数。确保文件夹包含子文件夹 val。

分布式训练

此代码已在多达 1024 个 A100 GPU 上经过实战测试。随着设备数量增加，logit 矩阵的空间复杂度会上升。使用全聚集方案复杂度为 O(n^2)，而使用 --gather-with-grad 和 --local-loss 标志可使复杂度变为有效线性。

多节点与 SLURM

对于大规模集群，推荐使用 SLURM 脚本。以下是一个训练最大模型的示例配置：

#!/bin/bash -x
#SBATCH --nodes=32
#SBATCH --gres=gpu:4
#SBATCH --ntasks-per-node=4
#SBATCH --cpus-per-task=6
#SBATCH --wait-all-nodes=1
#SBATCH --job-name=open_clip
#SBATCH --account=ACCOUNT_NAME
#SBATCH --partition PARTITION_NAME

eval "$(/path/to/conda/bin/conda shell.bash hook)"
conda activate open_clip
export CUDA_VISIBLE_DEVICES=0,1,2,3
export MASTER_PORT=12802
master_addr=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_ADDR=$master_addr
cd /shared/open_clip
export PYTHONPATH="$PYTHONPATH:$PWD/src"

srun --cpu_bind=v --accel-bind=gn python -u src/open_clip_train/main.py \
  --save-frequency 1 \
  --report-to tensorboard \
  --train-data="/data/LAION-400M/{00000..41455}.tar" \
  --warmup 2000 \
  --batch-size=256 \
  --epochs=32 \
  --workers=8 \
  --model ViT-B-32 \
  --name "ViT-B-32-Vanilla" \
  --seed 0 \
  --local-loss \
  --gather-with-grad

数据策略

OpenCLIP 支持通过 :: 分隔符使用多个数据源。例如：

--train-data "/data/cc12m/cc12m-train-{0000..2175}.tar::/data/LAION-400M/{00000..41455}.tar"

在这种情况下，建议使用 --dataset-resampled 进行替换采样。默认情况下，模型看到每个来源样本的次数与来源大小成正比。如果需要调整采样频率，可使用 --train-data-upsampling-factors 标志。

CoCa 模型训练

CoCa 模型训练通过指定 --model 参数启用。当前可用配置包括 coca_base、coca_ViT-B-32 等。CoCa 配置包含额外的 multimodal_cfg 组件。

微调 CoCa 时，关键是将对比损失权重设为 0，仅训练生成端：

--coca-contrastive-loss-weight 0 \
--coca-caption-loss-weight 1

评估与高级特性

零样本评估

系统评估建议在 40 个数据集上进行。评估本地检查点的命令如下：

python -m open_clip_train.main \
  --val-data="/path/to/validation_data.csv" \
  --model RN101 \
  --pretrained /path/to/checkpoints/epoch_K.pt

模型蒸馏

可以使用 --distill-model 和 --distill-pretrained 从预训练模型中提取知识。例如从 OpenAI ViT-L/14 蒸馏：

--distill-model ViT-L-14 --distill-pretrained openai

梯度累积

模拟更大的批次大小可使用 --accum-freq k。如果每 GPU 批次大小为 m，则有效批次大小为 k * m * num_gpus。增加此值前，建议先使用 --grad-checkpointing、--local-loss 等功能减少内存占用。

Int8 支持

目前提供 Int8 训练和推理的测试版支持。启用命令：

--use-bnb-linear SwitchBackLinearGlobal

对于 CLIP VIT-Huge，这通常能带来约 10% 的训练加速且无精度损失。

远程同步

支持直接从 S3 等远程文件系统恢复检查点，或在训练期间持续备份到 S3：

--logs /scratch \
--remote-sync s3://<path-to-bucket>

如需避免本地保存过多检查点，可添加 --delete-previous-checkpoint。

推送至 Hugging Face Hub

使用模块 open_clip.push_to_hf_hub 可将模型权重和配置推送到 HF Hub：

python -m open_clip.push_to_hf_hub \
  --model convnext_large_d_320 \
  --pretrained /train/checkpoints/epoch_12.pt \
  --repo-id laion/CLIP-convnext_large_d_320.laion2B-s29B-b131K-ft

致谢

感谢 Gauss Centre for Supercomputing e.V. 提供的计算时间支持。本仓库的开发由 Gabriel Ilharco 等人领导，原始版本来自华盛顿大学、谷歌、斯坦福等多机构研究人员合作。