前言
CLIP(Contrastive Language-Image Pre-training)作为OpenAI提出的跨模态对比学习经典模型,实现了图像-文本的双向语义对齐,凭借零样本迁移能力成为计算机视觉和自然语言处理跨模态任务的基础。本文将逐行解析CLIP的PyTorch原生实现源码,从基础模块到整体架构,深入理解其视觉编码器、文本编码器和对比学习核心逻辑,同时掌握其中的经典改进技巧(如改进ResNet、注意力池化、轻量化Transformer等)。
本文解析的源码为CLIP官方简化版实现,剔除了工程化冗余代码,保留核心模型架构,适合深度学习开发者、算法工程师深入理解CLIP的底层实现原理。
一、源码整体结构
本次解析的CLIP源码由基础模块层 、编码器层 、主模型层 和工具函数层 四层构成,核心逻辑围绕视觉编码器 (改进ResNet/Vision Transformer)、文本编码器 (轻量化Transformer)和跨模态对比学习前向传播展开,整体类结构如下:
├── 基础模块:Bottleneck/AttentionPool2d/LayerNorm/QuickGELU/ResidualAttentionBlock
├── 视觉编码器:ModifiedResNet(改进ResNet)/VisionTransformer(视觉Transformer)
├── 文本编码器:Transformer(轻量化自注意力Transformer)
├── 主模型:CLIP(融合视觉/文本编码器,实现跨模态对比学习)
├── 工具函数:convert_weights(权重类型转换)/build_model(从权重构建模型)
源码基于PyTorch实现,依赖numpy和Python内置库,无其他第三方框架依赖,可直接复现。
二、基础模块解析
基础模块是构建CLIP编码器的核心组件,包含改进的ResNet瓶颈块、注意力池化、自定义层归一化、激活函数和自注意力残差块,是理解后续编码器的基础。
2.1 Bottleneck:改进的ResNet瓶颈块
CLIP对传统ResNet的Bottleneck块做了抗锯齿下采样 改进,是ModifiedResNet的核心组件,其expansion=4保持与传统ResNet一致(输出通道数为输入的4倍)。
python
class Bottleneck(nn.Module):
expansion = 4
def __init__(self, inplanes, planes, stride=1):
super().__init__()
# 1x1卷积降维
self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
self.bn1 = nn.BatchNorm2d(planes)
self.relu1 = nn.ReLU(inplace=True)
# 3x3卷积提取特征(恒等步长,无下采样)
self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes)
self.relu2 = nn.ReLU(inplace=True)
# 关键改进:步长>1时用平均池化下采样,替代传统3x3卷积的步长下采样(抗锯齿)
self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
# 1x1卷积升维
self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
self.relu3 = nn.ReLU(inplace=True)
# 下采样捷径:步长>1或通道不匹配时,通过"平均池化+1x1卷积"对齐维度
self.downsample = None
self.stride = stride
if stride > 1 or inplanes != planes * Bottleneck.expansion:
self.downsample = nn.Sequential(OrderedDict([
("-1", nn.AvgPool2d(stride)),
("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
("1", nn.BatchNorm2d(planes * self.expansion))
]))
def forward(self, x):
identity = x
# 前向传播:1x1->3x3->AvgPool(可选)->1x1
out = self.relu1(self.bn1(self.conv1(x)))
out = self.relu2(self.bn2(self.conv2(out)))
out = self.avgpool(out)
out = self.bn3(self.conv3(out))
# 捷径连接:维度对齐后相加
if self.downsample is not None:
identity = self.downsample(x)
out += identity
out = self.relu3(out)
return out
核心改进点:
- 传统ResNet通过3x3卷积的
stride>1实现下采样,CLIP则将所有卷积步长设为1,步长>1时通过平均池化实现下采样,避免卷积下采样的锯齿效应,提升特征提取的平滑性; - 捷径连接的下采样同样采用"平均池化+1x1卷积"的组合,而非直接用1x1卷积下采样,保证维度对齐的同时保持特征平滑。
2.2 AttentionPool2d:注意力池化层
替代传统ResNet的全局平均池化(GAP),是CLIP视觉编码器的关键创新,通过QKV自注意力对视觉特征进行池化,能够自适应捕捉特征图的全局语义信息,而非简单的均值求和。
python
class AttentionPool2d(nn.Module):
def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
super().__init__()
# 位置嵌入:(H*W+1)×embed_dim,+1为全局均值特征的位置
self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
# QKV投影层
self.k_proj = nn.Linear(embed_dim, embed_dim)
self.q_proj = nn.Linear(embed_dim, embed_dim)
self.v_proj = nn.Linear(embed_dim, embed_dim)
# 输出投影层
self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
self.num_heads = num_heads
def forward(self, x):
# 特征重塑:NCHW -> (HW)NC(将空间维度展平,放到第一维)
x = x.flatten(start_dim=2).permute(2, 0, 1)
# 拼接全局均值特征:(HW)NC -> (HW+1)NC,第一个特征为全局均值,作为查询(Query)的基础
x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)
# 加入位置嵌入
x = x + self.positional_embedding[:, None, :].to(x.dtype)
# 多头自注意力前向传播:仅用第一个特征(全局均值)作为Query,所有特征作为Key/Value
x, _ = F.multi_head_attention_forward(
query=x[:1], key=x, value=x,
embed_dim_to_check=x.shape[-1], num_heads=self.num_heads,
q_proj_weight=self.q_proj.weight, k_proj_weight=self.k_proj.weight, v_proj_weight=self.v_proj.weight,
in_proj_weight=None, in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
out_proj_weight=self.c_proj.weight, out_proj_bias=self.c_proj.bias,
use_separate_proj_weight=True, training=self.training, need_weights=False
)
return x.squeeze(0)
核心逻辑:
- 特征重塑:将卷积输出的
NCHW格式展平为(HW)NC,把每个空间位置的特征作为一个"token"; - 全局均值token:计算所有空间token的均值,作为查询(Query) 基础,拼接在原有token前,形成
HW+1个token; - 多头自注意力:仅用全局均值token作为Query,所有token(含均值)作为Key和Value,实现全局语义的自适应聚合;
- 输出:注意力输出后挤压维度,得到
N×output_dim的视觉特征向量。
2.3 LayerNorm & QuickGELU:自定义归一化和激活函数
为适配半精度(FP16)训练和提升模型收敛速度,CLIP实现了自定义的LayerNorm和轻量化GELU:
python
class LayerNorm(nn.LayerNorm):
"""适配FP16的LayerNorm,先转FP32计算,再转回原类型"""
def forward(self, x: torch.Tensor):
orig_type = x.dtype
ret = super().forward(x.type(torch.float32))
return ret.type(orig_type)
class QuickGELU(nn.Module):
"""轻量化GELU,用x*sigmoid(1.702x)近似原GELU,减少计算量"""
def forward(self, x: torch.Tensor):
return x * torch.sigmoid(1.702 * x)
设计初衷:
- 原生LayerNorm在FP16下易出现数值不稳定,自定义版本先转换为FP32完成计算,再转回原数据类型,保证数值稳定性;
- 原GELU函数包含
erf操作,计算量较大,QuickGELU通过sigmoid近似实现,在精度损失极小的前提下提升计算效率,是CLIP和Stable Diffusion等模型的常用技巧。
2.4 ResidualAttentionBlock:自注意力残差块
构建CLIP Transformer的基础单元,融合多头自注意力 和前馈网络(MLP),采用经典的"预归一化+残差连接"架构(与GPT一致),是视觉Transformer和文本Transformer的通用组件。
python
class ResidualAttentionBlock(nn.Module):
def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
super().__init__()
self.attn = nn.MultiheadAttention(d_model, n_head) # 多头自注意力
self.ln_1 = LayerNorm(d_model) # 自注意力前的层归一化
# 前馈网络:4倍扩维 -> QuickGELU -> 原维数投影
self.mlp = nn.Sequential(OrderedDict([
("c_fc", nn.Linear(d_model, d_model * 4)),
("gelu", QuickGELU()),
("c_proj", nn.Linear(d_model * 4, d_model))
]))
self.ln_2 = LayerNorm(d_model) # MLP前的层归一化
self.attn_mask = attn_mask # 注意力掩码(文本Transformer用因果掩码,视觉用None)
def attention(self, x: torch.Tensor):
# 掩码适配:保证掩码的类型和设备与输入一致
self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
def forward(self, x: torch.Tensor):
# 自注意力残差连接:x = x + Attn(LN(x))
x = x + self.attention(self.ln_1(x))
# MLP残差连接:x = x + MLP(LN(x))
x = x + self.mlp(self.ln_2(x))
return x
核心架构:
- 预归一化:与传统Transformer的"后归一化"不同,先对输入做LayerNorm再送入自注意力/MLP,提升模型收敛速度和稳定性;
- 残差连接:自注意力和MLP的输出均与原始输入相加,缓解深度网络的梯度消失问题;
- 注意力掩码:支持自定义掩码,文本Transformer中用于因果掩码(防止看到未来的token),视觉Transformer中无需掩码。
三、核心编码器解析
CLIP的核心是双编码器架构 :视觉编码器(支持改进ResNet和Vision Transformer)和文本编码器(轻量化Transformer),两者均将输入映射到同一维度的语义嵌入空间,实现跨模态对齐。
3.1 Transformer:轻量化通用Transformer
基于ResidualAttentionBlock构建,是视觉Transformer(VisionTransformer类)和文本编码器的通用基础,结构简洁,无额外冗余模块。
python
class Transformer(nn.Module):
def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None):
super().__init__()
self.width = width # 特征维度
self.layers = layers # 自注意力块数量
# 堆叠layers个ResidualAttentionBlock
self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)])
def forward(self, x: torch.Tensor):
return self.resblocks(x)
设计特点:
- 纯自注意力堆叠:无编码器-解码器结构,仅由多层自注意力残差块构成,属于仅编码器(Encoder-only) 架构;
- 输入格式:要求输入为
L×N×D格式(L:序列长度,N:批次大小,D:特征维度),与PyTorch的MultiheadAttention输入格式一致; - 通用性:通过不同的
attn_mask适配视觉(无掩码)和文本(因果掩码)场景。
3.2 ModifiedResNet:CLIP改进版ResNet
CLIP对传统ResNet做了三处关键改进,使其更适合跨模态对比学习,替代原生ResNet作为视觉编码器之一。
python
class ModifiedResNet(nn.Module):
def __init__(self, layers, output_dim, heads, input_resolution=224, width=64):
super().__init__()
self.output_dim = output_dim # 最终输出嵌入维度
self.input_resolution = input_resolution # 输入图像分辨率
# 改进点1:3层Stem卷积(替代传统1层)+ 平均池化(替代最大池化)
self.conv1 = nn.Conv2d(3, width // 2, 3, stride=2, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(width // 2)
self.relu1 = nn.ReLU(inplace=True)
self.conv2 = nn.Conv2d(width // 2, width // 2, 3, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(width // 2)
self.relu2 = nn.ReLU(inplace=True)
self.conv3 = nn.Conv2d(width // 2, width, 3, padding=1, bias=False)
self.bn3 = nn.BatchNorm2d(width)
self.relu3 = nn.ReLU(inplace=True)
self.avgpool = nn.AvgPool2d(2)
# 残差层:基于改进的Bottleneck块,步长>1时用平均池化下采样
self._inplanes = width
self.layer1 = self._make_layer(width, layers[0])
self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
# 改进点3:注意力池化(AttentionPool2d)替代全局平均池化
embed_dim = width * 32 # ResNet最终特征维度
self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, output_dim)
def _make_layer(self, planes, blocks, stride=1):
# 构建残差层:第一个块做下采样,后续块恒等映射
layers = [Bottleneck(self._inplanes, planes, stride)]
self._inplanes = planes * Bottleneck.expansion
for _ in range(1, blocks):
layers.append(Bottleneck(self._inplanes, planes))
return nn.Sequential(*layers)
def forward(self, x):
# Stem层前向传播
def stem(x):
x = self.relu1(self.bn1(self.conv1(x)))
x = self.relu2(self.bn2(self.conv2(x)))
x = self.relu3(self.bn3(self.conv3(x)))
x = self.avgpool(x)
return x
# 适配权重数据类型(如FP16)
x = x.type(self.conv1.weight.dtype)
x = stem(x)
# 残差层前向传播
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
# 注意力池化得到最终视觉特征
x = self.attnpool(x)
return x
CLIP对ResNet的三大核心改进:
- Stem层升级:将传统ResNet的1个7x7卷积(步长2)改为3个3x3卷积(第一个步长2,后两个步长1),并将最大池化改为平均池化,减少大卷积核的信息损失,提升特征提取的精细度;
- 抗锯齿下采样 :通过
Bottleneck块中的平均池化实现下采样,替代卷积步长下采样,避免锯齿效应; - 注意力池化替代GAP :用
AttentionPool2d替代传统的全局平均池化,自适应聚合全局语义特征,而非简单均值,提升特征的表达能力。
3.3 VisionTransformer:CLIP版视觉Transformer
CLIP的另一款视觉编码器,采用经典的ViT架构,做了少量适配跨模态学习的优化,与ModifiedResNet二选一使用。
python
class VisionTransformer(nn.Module):
def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int):
super().__init__()
self.input_resolution = input_resolution # 输入分辨率
self.output_dim = output_dim # 输出嵌入维度
# 图像分块嵌入:3x3卷积(步长=块大小)替代展平+线性层,实现Patch Embedding
self.conv1 = nn.Conv2d(3, width, kernel_size=patch_size, stride=patch_size, bias=False)
scale = width ** -0.5
# 类别嵌入:新增一个可学习token,与ViT一致,作为全局特征聚合的载体
self.class_embedding = nn.Parameter(scale * torch.randn(width))
# 位置嵌入:(num_patches+1)×width,+1为类别嵌入
self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width))
self.ln_pre = LayerNorm(width) # Transformer前的层归一化
# 视觉Transformer主体
self.transformer = Transformer(width, layers, heads)
# 输出层:层归一化 + 线性投影(映射到跨模态嵌入空间)
self.ln_post = LayerNorm(width)
self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
def forward(self, x: torch.Tensor):
# 1. Patch Embedding:NCHW -> N×width×grid×grid(grid=分辨率/块大小)
x = self.conv1(x)
# 2. 特征重塑:N×width×grid² -> N×grid²×width
x = x.reshape(x.shape[0], x.shape[1], -1).permute(0, 2, 1)
# 3. 拼接类别嵌入:N×(grid²+1)×width
x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)
# 4. 加入位置嵌入
x = x + self.positional_embedding.to(x.dtype)
# 5. 预归一化
x = self.ln_pre(x)
# 6. Transformer前向:NLD -> LND -> NLD(适配MultiheadAttention输入格式)
x = x.permute(1, 0, 2)
x = self.transformer(x)
x = x.permute(1, 0, 2)
# 7. 提取类别嵌入的特征并做后归一化
x = self.ln_post(x[:, 0, :])
# 8. 投影到跨模态嵌入空间
if self.proj is not None:
x = x @ self.proj
return x
CLIP-ViT的核心特点:
- Patch Embedding优化:用卷积层(核大小=步长=块大小)实现图像分块,替代传统的"展平+线性层",计算更高效,且能利用卷积的硬件加速;
- 类别嵌入 :新增一个可学习的
class_embedding,拼接在所有Patch嵌入前,作为全局语义特征的载体,最终仅提取该token的特征作为视觉输出; - 预/后归一化 :Transformer前做
ln_pre,提取类别特征后做ln_post,保证特征分布的稳定性; - 线性投影 :通过可学习的
proj将Transformer的输出特征映射到跨模态嵌入空间,与文本编码器的输出维度一致。
四、CLIP主模型解析
CLIP类是整个模型的入口,融合视觉编码器 和文本编码器 ,实现图像/文本特征编码 和跨模态对比学习的前向传播,核心是将图像和文本映射到同一嵌入空间,并通过余弦相似度计算跨模态匹配分数。
4.1 初始化与参数初始化
python
class CLIP(nn.Module):
def __init__(self,
embed_dim: int, # 跨模态嵌入空间维度
# 视觉参数
image_resolution: int, vision_layers: Union[Tuple[int, int, int, int], int],
vision_width: int, vision_patch_size: int,
# 文本参数
context_length: int, vocab_size: int,
transformer_width: int, transformer_heads: int, transformer_layers: int
):
super().__init__()
self.context_length = context_length # 文本序列最大长度
# 选择视觉编码器:元组类型layers->ModifiedResNet,整数->VisionTransformer
if isinstance(vision_layers, (tuple, list)):
vision_heads = vision_width * 32 // 64
self.visual = ModifiedResNet(layers=vision_layers, output_dim=embed_dim, heads=vision_heads,
input_resolution=image_resolution, width=vision_width)
else:
vision_heads = vision_width // 64
self.visual = VisionTransformer(input_resolution=image_resolution, patch_size=vision_patch_size,
width=vision_width, layers=vision_layers, heads=vision_heads, output_dim=embed_dim)
# 构建文本Transformer(带因果掩码)
self.transformer = Transformer(width=transformer_width, layers=transformer_layers,
heads=transformer_heads, attn_mask=self.build_attention_mask())
# 文本嵌入层:词嵌入 + 位置嵌入
self.vocab_size = vocab_size
self.token_embedding = nn.Embedding(vocab_size, transformer_width)
self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width))
self.ln_final = LayerNorm(transformer_width)
# 文本特征投影层:映射到跨模态嵌入空间
self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
# 温度系数:可学习,用于缩放余弦相似度
self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
# 初始化所有可学习参数
self.initialize_parameters()
关键初始化逻辑:
- 双视觉编码器适配 :通过
vision_layers的类型判断使用ModifiedResNet(元组/列表,对应四层残差块的数量)还是VisionTransformer(整数,对应Transformer层数); - 文本因果掩码 :通过
build_attention_mask构建上三角掩码,实现文本的自回归注意力(防止看到未来的token); - 温度系数
logit_scale:初始值为np.log(1/0.07),通过指数运算后作为余弦相似度的缩放系数,是对比学习的关键超参数,设为可学习参数让模型自适应调整; - 跨模态投影 :文本编码器通过
text_projection、视觉编码器通过内部的投影层,将各自的特征映射到同一embed_dim维度的嵌入空间。
4.2 自定义参数初始化
CLIP针对不同模块设计了差异化的参数初始化策略,保证模型初始状态的数值稳定性和收敛速度。
python
def initialize_parameters(self):
# 词嵌入和位置嵌入:正态分布初始化
nn.init.normal_(self.token_embedding.weight, std=0.02)
nn.init.normal_(self.positional_embedding, std=0.01)
# ModifiedResNet的注意力池化层初始化
if isinstance(self.visual, ModifiedResNet):
if self.visual.attnpool is not None:
std = self.visual.attnpool.c_proj.in_features ** -0.5
nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std)
nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std)
nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std)
nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std)
# ResNet瓶颈块的bn3.weight初始化为0:让残差连接初始时占主导
for resnet_block in [self.visual.layer1, self.visual.layer2, self.visual.layer3, self.visual.layer4]:
for name, param in resnet_block.named_parameters():
if name.endswith("bn3.weight"):
nn.init.zeros_(param)
# Transformer参数初始化:自注意力和MLP分别设置不同的标准差
proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
attn_std = self.transformer.width ** -0.5
fc_std = (2 * self.transformer.width) ** -0.5
for block in self.transformer.resblocks:
nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
# 文本投影层初始化
if self.text_projection is not None:
nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
核心初始化技巧:
- ResNet的
bn3.weight初始化为0:让瓶颈块的初始输出近似为0,残差连接占主导,缓解深度网络的初始化梯度问题; - Transformer差异化初始化:自注意力层、输出投影层、MLP的扩维/投影层分别设置不同的正态分布标准差,适配各模块的计算特性;
- 注意力池化层:按输出维度的平方根倒数设置初始化标准差,保证初始注意力分数的数值稳定性。
4.3 因果注意力掩码构建
为文本Transformer构建上三角因果掩码,防止模型在处理文本序列时看到未来的token,保证自回归特性。
python
def build_attention_mask(self):
# 构建加法型注意力掩码,无效位置填充-∞
mask = torch.empty(self.context_length, self.context_length)
mask.fill_(float("-inf"))
mask.triu_(1) # 上三角部分保留-∞,下三角(含对角线)置0
return mask
掩码逻辑:
- PyTorch的
MultiheadAttention采用加法型掩码 ,掩码值为-inf的位置会在softmax后概率趋近于0,实现注意力屏蔽; triu_(1)将矩阵的上三角部分(行索引<列索引)设为-inf,即每个token仅能关注自身和前面的token,无法关注后面的token(未来token)。
4.4 图像/文本特征编码
CLIP提供两个独立的编码方法,分别将图像和文本转换为归一化前的跨模态特征向量,是零样本推理的核心接口。
python
@property
def dtype(self):
# 获取模型权重的默认数据类型(如FP16/FP32)
return self.visual.conv1.weight.dtype
def encode_image(self, image):
# 图像编码:适配数据类型后送入视觉编码器
return self.visual(image.type(self.dtype))
def encode_text(self, text):
# 文本编码:词嵌入->位置嵌入->Transformer->提取EOT特征->投影
x = self.token_embedding(text).type(self.dtype) # [batch_size, n_ctx, d_model]
x = x + self.positional_embedding.type(self.dtype)
x = x.permute(1, 0, 2) # NLD -> LND(适配Transformer输入)
x = self.transformer(x)
x = x.permute(1, 0, 2) # LND -> NLD
x = self.ln_final(x).type(self.dtype)
# 提取EOT(End of Text)token的特征:text.argmax(dim=-1)找到每个序列的最后一个有效token
x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
return x
文本编码关键细节:
- EOT特征提取 :CLIP的文本token中,最大索引的token为结束符(EOT) ,通过
text.argmax(dim=-1)找到每个文本序列的EOT位置,仅提取该位置的特征作为文本的全局语义特征; - 数据类型适配:所有操作均转换为模型权重的默认类型,保证数值一致性。
4.5 前向传播:跨模态对比学习核心
CLIP的前向传播实现了对比学习的损失计算前逻辑 ,将图像和文本特征归一化后,计算图像-文本的余弦相似度矩阵,并通过温度系数缩放,得到匹配分数。
python
def forward(self, image, text):
# 1. 编码得到图像和文本特征
image_features = self.encode_image(image)
text_features = self.encode_text(text)
# 2. 特征归一化:L2归一化,保证余弦相似度的取值范围为[-1,1]
image_features = image_features / image_features.norm(dim=1, keepdim=True)
text_features = text_features / text_features.norm(dim=1, keepdim=True)
# 3. 温度系数指数化:将对数空间的系数转换为原始空间
logit_scale = self.logit_scale.exp()
# 4. 计算余弦相似度矩阵并缩放
logits_per_image = logit_scale * image_features @ text_features.t() # 图像到文本的匹配分数
logits_per_text = logits_per_image.t() # 文本到图像的匹配分数
# 输出:[batch_size, batch_size]的相似度矩阵
return logits_per_image, logits_per_text
跨模态对比学习核心逻辑:
- L2归一化 :对图像和文本特征做L2归一化后,两个特征的点积等价于余弦相似度,简化计算;
- 相似度矩阵 :
logits_per_image[i][j]表示第i张图像与第j个文本的匹配分数,对角线为正样本对(图像-文本匹配),非对角线为负样本对; - 温度系数缩放:缩放余弦相似度,让模型的匹配分数更适合通过交叉熵损失训练(温度系数越大,相似度分布越平缓,反之越陡峭)。
训练损失 :CLIP的训练损失为双向交叉熵损失 ,即对logits_per_image做行方向的交叉熵(图像匹配文本),对logits_per_text做行方向的交叉熵(文本匹配图像),最终取两者的均值。
五、工具函数解析
源码提供两个实用工具函数,分别实现模型权重的半精度转换 和从预训练权重构建CLIP模型,是CLIP工程化部署和预训练模型加载的核心。
5.1 convert_weights:权重半精度转换
将模型的卷积层、线性层、自注意力层等参数转换为FP16半精度,减少模型显存占用,提升推理速度,适配GPU推理场景。
python
def convert_weights(model: nn.Module):
def _convert_weights_to_fp16(l):
# 卷积层、线性层:权重和偏置转FP16
if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
l.weight.data = l.weight.data.half()
if l.bias is not None:
l.bias.data = l.bias.data.half()
# 多头自注意力层:所有投影权重和偏置转FP16
if isinstance(l, nn.MultiheadAttention):
for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
tensor = getattr(l, attr)
if tensor is not None:
tensor.data = tensor.data.half()
# 跨模态投影层:转FP16
for name in ["text_projection", "proj"]:
if hasattr(l, name):
attr = getattr(l, name)
if attr is not None:
attr.data = attr.data.half()
# 递归应用权重转换
model.apply(_convert_weights_to_fp16)
转换原则 :仅转换可学习的参数层,归一化层(BatchNorm/LayerNorm)不做转换,避免数值不稳定。
5.2 build_model:从预训练权重构建模型
CLIP的核心工程化函数,无需手动指定模型参数,通过解析预训练权重的键值对,自动推断模型的结构参数(如视觉编码器类型、层数、特征维度等),并构建模型、加载权重、转换为半精度。
python
def build_model(state_dict: dict):
# 判断视觉编码器类型:含visual.proj则为VisionTransformer,否则为ModifiedResNet
vit = "visual.proj" in state_dict
if vit:
# 解析VisionTransformer参数
vision_width = state_dict["visual.conv1.weight"].shape[0]
vision_layers = len([k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
image_resolution = vision_patch_size * grid_size
else:
# 解析ModifiedResNet参数
counts: list = [len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1,2,3,4]]
vision_layers = tuple(counts)
vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
vision_patch_size = None
assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
image_resolution = output_width * 32
# 解析跨模态和文本编码器参数
embed_dim = state_dict["text_projection"].shape[1]
context_length = state_dict["positional_embedding"].shape[0]
vocab_size = state_dict["token_embedding.weight"].shape[0]
transformer_width = state_dict["ln_final.weight"].shape[0]
transformer_heads = transformer_width // 64
transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith("transformer.resblocks")))
# 构建CLIP模型
model = CLIP(
embed_dim, image_resolution, vision_layers, vision_width, vision_patch_size,
context_length, vocab_size, transformer_width, transformer_heads, transformer_layers
)
# 删除状态字典中无关的键(避免加载权重时出错)
for key in ["input_resolution", "context_length", "vocab_size"]:
if key in state_dict:
del state_dict[key]
# 转换为FP16并加载预训练权重
convert_weights(model)
model.load_state_dict(state_dict)
# 设置为评估模式(关闭Dropout/BatchNorm的训练模式)
return model.eval()
核心优势:
- 自动参数推断:无需用户记忆CLIP的各版本参数(如ViT-B/32、ResNet50),通过预训练权重自动推断,降低使用门槛;
- 权重兼容性:删除状态字典中与模型结构无关的键,避免加载权重时的键值不匹配错误;
- 一键部署:自动转换为FP16并设置为评估模式,加载后可直接用于推理。
六、CLIP模型的使用流程
基于上述源码,CLIP模型的预训练权重加载 和零样本推理流程非常简洁,核心步骤如下:
python
import torch
import clip
from PIL import Image
# 1. 加载预训练权重(需提前下载CLIP预训练权重文件,如clip_vit_b32.pth)
state_dict = torch.load("clip_vit_b32.pth", map_location="cuda")
model = build_model(state_dict)
model = model.cuda() # 移至GPU
# 2. 图像预处理(需与CLIP训练时的预处理一致:224x224、归一化等)
image = Image.open("cat.jpg").convert("RGB")
preprocess = clip.load_default_preprocess() # 自定义或使用CLIP默认预处理
image_tensor = preprocess(image).unsqueeze(0).cuda()
# 3. 文本预处理(分词、填充到context_length)
texts = ["a photo of a cat", "a photo of a dog", "a photo of a bird"]
text_tokens = clip.tokenize(texts).cuda() # 自定义分词器或使用CLIP默认分词器
# 4. 零样本推理
with torch.no_grad():
image_features = model.encode_image(image_tensor)
text_features = model.encode_text(text_tokens)
# 计算匹配分数
image_features = image_features / image_features.norm(dim=1, keepdim=True)
text_features = text_features / text_features.norm(dim=1, keepdim=True)
logits = model.logit_scale.exp() * image_features @ text_features.t()
# 预测结果
pred = logits.argmax(dim=1).item()
print(f"预测结果:{texts[pred]}")
关键注意点:
- 图像和文本的预处理必须与CLIP训练时一致,否则会严重影响模型性能;
- 推理时需使用
torch.no_grad()关闭梯度计算,减少显存占用; - 文本分词需保证序列长度不超过
model.context_length,并做填充/截断处理。
七、总结与核心亮点回顾
本文逐行解析了CLIP的PyTorch原生实现源码,从基础模块到整体架构,深入理解了其跨模态对比学习的核心原理。CLIP的成功不仅源于图像-文本的对比学习范式,更得益于其精心设计的模型架构,核心亮点可总结为:
- 双编码器架构:视觉编码器(改进ResNet/ViT)和文本编码器(轻量化Transformer)将输入映射到同一嵌入空间,实现跨模态语义对齐;
- 视觉编码器创新:改进ResNet的抗锯齿下采样、注意力池化,ViT的卷积式Patch Embedding,提升视觉特征的表达能力;
- 高效的Transformer设计:预归一化、QuickGELU、因果掩码,兼顾模型收敛速度和计算效率;
- 对比学习优化:可学习的温度系数、双向交叉熵损失、L2归一化,让模型快速收敛到跨模态嵌入空间;
- 工程化友好:自动权重解析、半精度转换、模块化设计,降低模型使用和部署门槛。
CLIP作为跨模态学习的里程碑模型,其源码设计兼具学术创新性 和工程实用性,是深度学习开发者学习跨模态学习、Transformer、对比学习的绝佳范例。基于CLIP的思想,后续衍生出BLIP、ALBEF、FLAVA等一系列跨模态模型,推动了计算机视觉和自然语言处理的融合发展。
源码拓展方向:
- 基于该源码实现CLIP的训练代码,适配自定义的图像-文本数据集;
- 结合LoRA/QLoRA对CLIP进行微调,提升特定下游任务的零样本性能;
- 将CLIP的视觉编码器与扩散模型结合,实现文本到图像的生成(如Stable Diffusion)。
本文解析的源码为CLIP官方简化版,完整工程化代码可参考OpenAI官方CLIP仓库:https://github.com/openai/CLIP