ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

PyTorch实现CLIP模型:从双塔结构到零样本图像分类

PyTorch实现CLIP模型:从双塔结构到零样本图像分类 简介基于PyTorch实现的CLIP模型完整项目面向深度学习开发者、CV/NLP领域研究者及入门学习者解决图像与文本跨模态语义对齐与理解问题涵盖从数据准备到模型推理的完整链路。资源包共13个文件包含7个Python脚本覆盖模型架构、数据加载、训练与推理、1个Jupyter Notebook流程教程、1份README说明文档及3张示例图片整体仅3.11MB轻量易部署。目前已有817人学习下载。源码按PyTorch标准工程组织模块划分清晰涵盖模型定义、数据处理、训练、验证与推理等环节配套Notebook和Markdown教程深入讲解对比学习预训练原理、数据配对策略与训练参数调优方法可帮助用户快速复现图像分类、图文检索、视觉问答等任务。项目强调实战代码注释清晰易于修改与扩展也适合作为大模型相关课程设计与毕业设计的起点。1. 一个例子看清CLIP在解决什么没有标注数据也能给图片分类接到一个活要把几千张汽车图按“轿车、SUV、皮卡”区分开但库里连一张标签都没有。传统做法是先花一周人工标注、训练分类器、调参上线周期长且换一批车型就得重来。CLIPContrastive Language-Image Pre-training的思路是让模型提前学会“图片和文字在同一个语义空间里对齐”到了新任务上直接拿类别名去匹配图片一个样本都不用标。这个标题给的是一个“简洁明了的CLIP模型”落地包Pytorch实现、双塔结构、对比损失、附带项目源码和流程教程。定位很实在——不是把OpenAI原版几十亿数据训练的流程搬过来而是把结构讲清楚、让普通机器能跑通、让新手能照着复现。下面我按自己做过的方案从原理、代码、训练参数到踩坑记录完整讲一遍怎么把这个项目在本地落地。2. 双塔结构图像侧和文本侧各自提取特征再对对齐2.1 对比学习让图片和文本在同一个向量空间里相遇CLIP的核心不是分类而是度量学习。训练数据是一批图片, 文本配对比如一张狗的照片配一句“a photo of a dog”。训练时模型把图片编码成一个向量把文本也编码成一个向量目标很直观配对的图文向量要靠近同一batch里其他不配对的要拉开。这个batch内拉近推远的做法叫InfoNCE也叫对比学习。它和分类的本质区别在于分类是模型学会“这张图是狗”CLIP是模型学会“狗这个概念的图片表达”和“狗这个概念的文本表达”落在同一个向量区域里。所以换到新任务时不需要重新训练只要把新类别的名字写成句子跟图片向量做相似度排序就能出分类结果这就是zero-shot分类。训练样本的组织方式也值得注意。CLIP没有人为造负样本它直接把当前batch里的其他图文对当作负样本。这意味着batch size越大负样本越丰富学出来的特征越有区分度。原论文用了几万的大batch普通机器做不到但结构不变、损失函数不变小batch也能学到可用的特征只是效果上限有差距。损失函数长这样对一个batch的图片特征和文本特征做矩阵乘法得到B×B的相似度矩阵对角线是正样本其他位置全是负样本。然后分别按行、按列做交叉熵两个方向平均。这个对称形式让模型既要“图片找文本”也要“文本找图片”两边都学。相似度计算前有一个关键操作特征向量要做L2归一化。不归一化向量的模长会干扰相似度排序训练容易震荡。归一化之后点积的范围被限制在[-1, 1]再乘一个可学习的温度因子控制分布的锐度。温度越小softmax越尖锐模型越敢下判断。2.2 为什么用Pytorch自己实现而不是直接调open_clip这个决定不少人犹豫过。open_clip有现成权重、有训练好的模型pip装完就能用。那为什么还要自己用Pytorch写一遍我列出三条实际考虑。第一黑匣子问题。open_clip的代码封装层级很多想改结构、想换backbone、想只看一个batch的中间特征都要翻很久源码。自己做项目时我经常需要打印中间层的shape、检查某一层的梯度自己写的结构一眼就能定位。标题里“简洁明了”这四个字对应的就是这种可控性。第二微调和部署需要改结构。比如要把图像编码器从ResNet换成ViT或者要在文本编码器后面接一个自己的分类头open_clip的接口不一定支持直接改。自己维护的代码改动成本最低。后面章节我会讲到微调时的冻结策略自己写的模型改起来非常顺手。第三Pytorch的生态配套太省事。torchvision自带预训练ResNet和ViTtorch.cuda.amp一行开混合精度torch.utils.data.DataLoader的num_workers、pin_memory直接加速这些在训练流程里都是刚需。换TensorFlow或JAX同等功能的代码量至少翻一倍。另外说一个反向理由自己实现不代表不能用别人的预训练权重。图像编码器直接用torchvision的ResNet50预训练权重文本编码器从头训练这种组合在小规模数据上效果不错。CLIP学的是跨模态对齐视觉特征已经有了模型只需要学会把文本也映射到同一个空间。3. 用Pytorch写一个极简CLIP模型定义与对比损失3.1 图像编码器用预训练ResNet做backbone丢掉分类头图像侧最常见的做法是拿torchvision里的ResNet50当backbone去掉最后一层全局池化和分类头。为什么选ResNet50而不是更大模型因为CLIP是双塔结构文本侧也要训练两个塔一起算梯度视觉backbone太大显存扛不住。先跑通流程再考虑换ViT。import torch import torch.nn as nn import torch.nn.functional as F from torchvision import models class ImageEncoder(nn.Module): def __init__(self, out_dim512): super().__init__() # 加载IMAGENET1K_V2的预训练权重 backbone models.resnet50(weightsmodels.ResNet50_Weights.IMAGENET1K_V2) # 丢掉ResNet最后的全局平均池化和全连接分类头 self.backbone nn.Sequential(*list(backbone.children())[:-1]) # ResNet50 conv5输出2048维投影到统一的语义维度 self.proj nn.Linear(2048, out_dim) def forward(self, images): # images: [B, 3, H, W]输入尺寸建议224x224 features self.backbone(images) # [B, 2048, 1, 1] features features.flatten(1) # [B, 2048] features self.proj(features) # [B, out_dim] return F.normalize(features, dim-1) # 归一化关键不能省这里有个细节值得说nn.Sequential(*list(backbone.children())[:-1])把ResNet最后一个stage的输出保留下来没有走全局平均池化直接flatten。好处是保留了空间信息虽然CLIP最后只用全局特征但如果你想在中间层接自己的池化策略这个结构更好改。投影层self.proj把2048维压到统一维度。out_dim在整个双塔里必须一致图像侧和文本侧都输出这个维度才能做点积。我用512显存紧张可以降到256特征区分度会略降。注意F.normalize在forward里做了后续计算相似度时不要再归一化一次否则等于白做。3.2 文本编码器一个能处理句子的TransformerEncoder文本侧比图像侧更讲究。原版CLIP用的是GPT-2风格的Transformer输入是BPE编码的token序列句子前加[SOS]、末尾加[EOS]取EOS位置的输出作为句向量。简洁版可以省掉[SOS]保留EOS位置逻辑用Pytorch自带的nn.TransformerEncoder搭一个。class TextEncoder(nn.Module): def __init__(self, vocab_size, max_len64, embed_dim512, out_dim512, nhead8, num_layers3): super().__init__() self.token_emb nn.Embedding(vocab_size, embed_dim) # 可学习位置编码比正弦编码好在不用处理长度外推 self.pos_emb nn.Parameter(torch.zeros(1, max_len, embed_dim)) layer nn.TransformerEncoderLayer( d_modelembed_dim, nheadnhead, dim_feedforward2048, dropout0.1, batch_firstTrue, activationgelu ) self.encoder nn.TransformerEncoder(layer, num_layersnum_layers) self.proj nn.Linear(embed_dim, out_dim) def forward(self, tokens, mask): # tokens: [B, L]mask为True的位置是padding不参与attention B, L tokens.shape x self.token_emb(tokens) self.pos_emb[:, :L, :] x self.encoder(x, src_key_padding_maskmask) # 取每个样本最后一个有效字符位置的输出作为句向量 lengths mask.sum(dim1, keepdimTrue) # [B, 1] last lengths.clamp(min1) - 1 x x[torch.arange(B, devicex.device), last.squeeze(1)] x self.proj(x) # [B, out_dim] return F.normalize(x, dim-1)两个参数值得细说。max_len控制句子长度上限原版CLIP是77实际训练里大部分标题不到20个token设64足够太长浪费显存。num_layers我选3TransformerEncoderLayer内部已经有自注意力前馈层归一化深度太大在几千条图文对的小数据上容易过拟合。句向量的取法我做了简化原版取EOS位置的输出我这里取最后一个非padding位置的输出。只要tokenize时在句子末尾统一追加一个EOS token这个位置就是EOS。省去了记录EOS索引的麻烦训练效果几乎一样。注意mask.sum(dim1)统计的是padding数量最后一个有效字符的索引是总长度 - padding数 - 1代码里clamp(min1)是为了防止空句子。3.3 对比损失对称的InfoNCE温度参数怎么设模型主体由两个塔组成损失函数是CLIP的灵魂。把两个塔的输出做矩阵乘法对角线是正样本对其余位置全是负样本。用交叉熵分别按行按列算一遍再平均这个对称损失让两个塔同步更新。class CLIPModel(nn.Module): def __init__(self, image_encoder, text_encoder, logit_scale_init2.659): super().__init__() self.image_encoder image_encoder self.text_encoder text_encoder # 温度参数的倒数初始化为 ln(1/0.07) ≈ 2.659 self.logit_scale nn.Parameter(torch.tensor(logit_scale_init)) def forward(self, images, tokens, mask): image_features self.image_encoder(images) # [B, D]已归一化 text_features self.text_encoder(tokens, mask) # [B, D] logits self.logit_scale * (image_features text_features.t()) labels torch.arange(logits.size(0), devicelogits.device) loss_img F.cross_entropy(logits, labels) # 图片去匹配文本 loss_txt F.cross_entropy(logits.t(), labels) # 文本去匹配图片 return (loss_img loss_txt) / 2温度参数这里我用的是logit_scale它是1/temperature。原论文初始化temperature0.07对应logit_scale约2.659。为什么不用temperature本身因为Pytorch里直接优化temperature会出现除零和梯度爆炸优化它的对数形式数值更稳定。logit_scale在训练中是可学习参数会自动调整相似度分布的锐度。训练初期logits普遍很小交叉熵loss高模型被迫拉大scale后期特征区分度好了scale会稳定在一个值附近。如果你发现训练结束后logit_scale异常大比如超过10说明特征没有学好模型在用温度硬撑这时候要先检查特征质量而不是调温度。4. 跑通训练和zero-shot推理数据处理、参数与评估4.1 数据组织一张CSV表同时管理图片和文本训练CLIP需要图文对数据。最省事的组织方式是一张CSV两列image_path和text。没有现成数据时常用起点是COCO的train2017图片加caption标注抽几千对就能跑通流程。关键是数据清洗一句文本对应一张图文本不能太长长句截断即可。import os import torch import pandas as pd from PIL import Image from torch.utils.data import Dataset class ImageTextDataset(Dataset): def __init__(self, csv_path, img_dir, transform, max_len64): self.df pd.read_csv(csv_path) self.img_dir img_dir self.transform transform self.max_len max_len def __len__(self): return len(self.df) def __getitem__(self, idx): row self.df.iloc[idx] image Image.open(os.path.join(self.img_dir, row[image_path])).convert(RGB) image self.transform(image) tokens, mask self.encode_text(row[text]) return image, tokens, mask def encode_text(self, text): # 用tiktoken的cl100k_base做BPE编码效果接近原版CLIP的tokenizer import tiktoken enc tiktoken.get_encoding(cl100k_base) ids enc.encode(text)[:self.max_len - 1] [enc.eot_token] length len(ids) tokens torch.zeros(self.max_len, dtypetorch.long) mask torch.ones(self.max_len, dtypetorch.bool) tokens[:length] torch.tensor(ids) mask[:length] False # False表示有效字符True表示padding return tokens, maskencode_text里的逻辑要和TextEncoder的forward对得上句子末尾追加eot_tokenmask里有效位置置False。这样做有两个好处一是Transformer的attention不会看到padding位置二是TextEncoder取“最后一个有效位置”时取到的正好是EOS的输出。tiktoken是OpenAI开源的BPE编码器cl100k_base词表约10万token。你也可以用HuggingFace的CLIPProcessor里的tokenizer效果等同。自己写tokenizer的坑在于词表不一致会导致Embedding维度对不上最省心的是训练和推理永远用同一个编码函数。transform部分图像侧缩放到224x224用torchvision的标准流程。注意图像增强不要加RandomCrop因为配对的文本描述的是整张图不是某个局部。可以用RandomResizedCrop配合小比例缩放但不要裁掉太多主体。from torchvision import transforms train_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p0.5), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])4.2 训练循环混合精度、学习率与日志训练配置上CLIP和普通分类模型的差别主要在两点一是batch size尽量大二是温度参数跟着学。先看训练循环的完整代码。device cuda if torch.cuda.is_available() else cpu enc tiktoken.get_encoding(cl100k_base) model CLIPModel( ImageEncoder(out_dim512), TextEncoder(vocab_sizeenc.n_vocab, max_len64, out_dim512) ).to(device) optimizer torch.optim.AdamW(model.parameters(), lr5e-5, weight_decay0.2) scaler torch.cuda.amp.GradScaler() model.train() for epoch in range(30): total_loss 0 for images, tokens, mask in dataloader: images images.to(device) tokens tokens.to(device) mask mask.to(device) optimizer.zero_grad() with torch.cuda.amp.autocast(): loss model(images, tokens, mask) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() total_loss loss.item() print(fepoch {epoch} loss {total_loss / len(dataloader):.4f})学习率5e-5配合AdamW是CLIP类模型的安全起点。weight_decay设0.2看起来偏高但对比学习任务里正则化能防止特征空间过度膨胀。如果你的数据只有几千对每个epoch不到100步30个epoch可能不够我习惯看到loss平原后把学习率降到1e-5再跑10个epoch。batch size如果上不了64用梯度累积模拟大batch这是对比学习任务里性价比最高的技巧。accum_steps 4 for step, (images, tokens, mask) in enumerate(dataloader): images images.to(device); tokens tokens.to(device); mask mask.to(device) with torch.cuda.amp.autocast(): loss model(images, tokens, mask) / accum_steps scaler.scale(loss).backward() if (step 1) % accum_steps 0: scaler.step(optimizer) scaler.update() optimizer.zero_grad()梯度累积的副作用是BN的统计量更新频率变低不过我们的图像编码器用的是预训练ResNetbackbone部分BN已经学好了影响不大。训练时监控三个指标loss是否下降、logit_scale是否上升后趋稳、每个epoch结束时在验证集上跑一次zero-shot精度。第三个指标才是CLIP有没有学好的最终标准。4.3 zero-shot推理把你的类别列表变成一句句模板训练完的CLIP怎么用不需要任何训练样本把类别名写成自然语言模板让文本编码器编码再和图片特征算相似度。CIFAR-10的类别模板就是“a photo of a {class}”class换成plane、car、bird等。def zeroshot_classify(model, images, class_names, enc, device): model.eval() texts [fa photo of a {c} for c in class_names] token_list, mask_list [], [] for t in texts: tokens, mask encode_text(t, enc) token_list.append(tokens); mask_list.append(mask) tokens torch.stack(token_list).to(device) masks torch.stack(mask_list).to(device) with torch.no_grad(): text_features model.text_encoder(tokens, masks) # [C, D] image_features model.image_encoder(images) # [B, D] similarity image_features text_features.t() # [B, C] preds similarity.argmax(dim-1) return preds模板的选择直接影响精度。直接写“car”效果差因为训练数据里文本基本都是完整句子模型没见过裸词。加“a photo of a”前缀是原论文验证过的最稳模板。特定数据集上可以更细化比如医疗图像用“a chest X-ray of {class}”这个细节后面微调章节还会提到。zero-shot推理时注意model.eval()会关掉Dropout影响不大但BN的running_mean在单卡推理时不更新这是预期行为。多次调用zeroshot_classify时相同类别文本会被重复编码可以在类外缓存text_features推理效率提升明显。5. 训练CLIP最容易踩的6个坑现象、原因与解法5.1 Loss变成nan训练直接中断现象训练到某一步loss变成nan之后无法恢复。原因有三个高发点logit_scale在AMP混合精度下溢出、文本编码器位置编码初始化过大、图像侧某个batch出现损坏数据。解决logit_scale初始化后加一行self.logit_scale.data.clamp_(max4.6)上限相当于温度0.217训练中不会超过这个值。位置编码用nn.init.normal_(self.pos_emb, std0.02)初始化。数据侧在Dataset的__getitem__里包一层try-except图片解码失败时返回同batch的上一张图避免PIL丢异常。5.2 loss在下降但zero-shot精度纹丝不动现象训练loss从8降到4换到分类任务上精度还是随机水平。原因图像侧用的是预训练ResNet文本侧从零开始模型发现靠图像特征已经能区分batch内的样本文本侧成了摆设根本没有学到语义对齐。解决前5个epoch冻结图像编码器只训练文本侧和投影层。具体做法是把image_encoder参数的requires_grad置False等loss明显下降后再解冻一起训练。这个技巧在图文数据量少于5万对时几乎是必须的。5.3 batch内图片和文本数量对不上现象DataLoader报错或者模型forward里矩阵乘法维度不匹配。原因文本长度不同虽然mask了padding但tensor的shape必须完全一致。我遇到过一个隐蔽版本encode_text里截断长度max_len设了64但tokens初始化用了torch.zeros(self.max_len)mask用了torch.ones(self.max_len)两个都是64没问题换了一个数据集重新实例化Dataset时max_len传了128DataLoader里旧数据的缓存没清报错信息指向模型内部排查了半天。解决训练前打印一个batch的shapefor batch in dataloader: print([x.shape for x in batch]); break。确认image是[B,3,224,224]、tokens和mask都是[B,64]再开始训练。5.4 GPU利用率低显存没满但训练很慢现象loss在降但nvidia-smi显示GPU利用率只有20%上下。原因DataLoader的num_workers默认是0每步都在主进程里等图片读取和解码。CLIP的数据加载比分类任务更重因为每次还要做BPE编码。解决DataLoader里设num_workers4, pin_memoryTrue, persistent_workersTrue。如果图片是原始大图先在预处理脚本里统一resize到256x256存成pngDataset里只做ToTensor和Normalize加载速度能快3倍。BPE编码不要在__getitem__里现算提前把所有文本编码成tokens和mask存成npy文件训练时直接查表。5.5 微调时发现图像编码器根本没更新现象loss下降极慢打印参数发现backbone的梯度全是None。原因torchvision的预训练模型加载后有人习惯性调了for param in backbone.parameters(): param.requires_grad False后面忘了解冻。另一个常见误操作是backbone子模块里的BN层处于eval模式导致running_mean不更新。解决训练前加一段检查代码total sum(1 for p in model.parameters() if p.requires_grad) print(ftrainable params: {total})再确认model.train()被调用。如果只想解冻部分层比如ResNet最后两个stage就只把那部分requires_grad置True前面保持False。5.6 单卡能跑通换多卡后loss明显变高现象同一份代码从单卡改DDPloss比单卡高而且不稳定。原因对比学习的负样本来自batch内部DDP下每个进程只看到自己的本地batch负样本数量变少学出来的特征区分度下降。这是在CLIP任务里多卡训练最常见的坑之一。解决小数据阶段用梯度累积替代多卡。如果一定要多卡需要在每个step用torch.distributed.all_gather把所有进程的图像特征和文本特征收集起来再算loss代码量增加不少。我的建议是先单卡把流程跑通数据量大了再投入精力做分布式。6. 进阶微调自己的CLIP并在CIFAR-10上做验证6.1 微调策略先冻结再解冻按数据分布决定解冻深度CLIP模型微调是热词实际场景里拿到一个预训练或自己训练的CLIP要在特定领域用常见做法是先冻结再解冻。数据分布跟自然图像接近比如电商商品图、社交媒体图片冻结backbone只训练投影层和文本侧就够学习率可以提到1e-45个epoch就能看到效果。数据分布差异大医学影像、卫星图、工业质检图至少解冻图像编码器最后两个stage让高层特征适配新域。解冻操作很简单找到model.image_encoder.backbone把最后两个BasicBlock或Bottleneck的requires_grad置True前面的保持False。优化器要用filter(lambda p: p.requires_grad, model.parameters())重新构造否则冻结层的梯度为None会报错。6.2 用CIFAR-10验证zero-shot效果跑通训练后第一步验证永远是zero-shot分类。CIFAR-10是标准基准类别是10个常见物体模板用“a photo of a {class}”。用torchvision加载测试集跑一遍zeroshot_classify记录准确率。检查项预期结果不达标时先查什么训练loss降到3.5以下检查梯度累积、学习率logit_scale稳定在4~6之间特征质量差检查数据清洗CIFAR-10 zero-shot55%~70%区间先换更强的文本模板单类准确率不要出现某类全错类别名和模板是否匹配ResNet50图像侧配合小规模图文对训练的CLIPCIFAR-10 zero-shot精度正常在55%~70%。原论文ViT-B/32能到76%差距主要在训练数据规模我们几千对数据学出来的特征区分度必然有限这个数值是健康的。我习惯在每个epoch训练结束后跑一次CIFAR-10验证不只为了看精度更是为了观察学的方向对不对。如果loss降了但精度没动说明两个塔在对齐一个错误的空间这时候先去看数据清洗和文本质量如果精度在涨但涨幅很慢优先检查温度参数和负样本数量。CLIP这个模型训练曲线的诡异程度比分类模型高得多不盯验证集很容易白跑十几个小时。希望帮到你。本文还有配套的精品资源点击获取
返回列表