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() 的输出兼容：

model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='/path/to/my/b32.pt')

若从 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

数据策略

使用 :: 分隔符可以指定多个数据源，例如：

--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 等，这些配置包含额外的 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 等人领导，原始版本来自华盛顿大学、谷歌、斯坦福等多机构研究人员合作。