|
# CLIP模型简介
CLIP（Contrastive Language-Image Pre-training）是OpenAI提出的多模态预训练模型，通过对比学习实现图像和文本的对齐。
### 核心思想
将图像和文本映射到同一特征空间，通过对比学习使匹配的图像-文本对在特征空间中更接近。
### 架构组成
```
输入图像 ─→ 图像编码器 ─→ 图像特征向量 (768维)
↓
共享特征空间
↓
输入文本 ─→ 文本编码器 ─→ 文本特征向量 (768维)
↓
相似度计算
```
### 依赖库
```python
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel
from PIL import Image
```
### 环境要求
- torch >= 2.0.0
- transformers >= 4.54.0
- CUDA支持（推荐）
|
|
# PMC-CLIP特点
### 医学领域专用
PMC-CLIP是专门针对医学图像-文本匹配的CLIP模型变体：
| 特点 | 说明 |
|------|------|
| **医学文本编码器** | 使用BiomedBERT而非通用BERT |
| **医学图像理解** | 针对X光、MRI等医学图像优化 |
| **开放获取数据** | 在PMC（PubMed Central）开放获取数据上训练 |
### 模型权重
可从Hugging Face下载预训练权重：
```bash
# 文本编码器（约418MB）
wget https://huggingface.co/datasets/axiong/pmc_oa/resolve/main/text_encoder.pth
# 图像编码器 ResNet50（约144MB）
wget https://huggingface.co/datasets/axiong/pmc_oa/resolve/main/image_encoder(resnet50).pth
# 文本投影层（约2.3MB）
wget https://huggingface.co/datasets/axiong/pmc_oa/resolve/main/text_projection_layer.pth
```
### 特征维度
- 图像特征：768维
- 文本特征：768维
- 对比温度参数（logit_scale）：4.4292
|
|
# 应用场景
### 医学图像-文本检索
| 场景 | 说明 |
|------|------|
| **报告生成** | 根据医学图像生成诊断报告 |
| **图像检索** | 通过文本描述检索相关医学图像 |
| **多模态检索** | 图像查询文本或文本查询图像 |
| **诊断辅助** | 辅助医生快速定位相似病例 |
### 典型任务
1. **零样本分类**：无需训练即可分类医学图像
2. **跨模态检索**：文本描述查找相应图像
3. **报告生成**：图像自动生成医学报告
4. **质量评估**：评估医学图像与报告的一致性
### 示例数据
- 胸部X光片（Chest X-ray）
- 脑部MRI（Brain MRI）
- CT扫描图像
- 病理切片图像
|
|
# Modified ResNet架构
### 核心改进
PMC-CLIP使用的Modified ResNet相比标准ResNet有以下改进：
1. **三层Stem**：替换原来的单层卷积
2. **抗混叠卷积**：在步长>1的卷积前添加平均池化
3. **注意力池化**：替换最后的全局平均池化
### 网络结构
```
输入图像 (224×224×3)
↓
Stem (3层卷积 + AvgPool)
↓
Layer1 (3个Bottleneck, 64维)
↓
Layer2 (4个Bottleneck, 128维, stride=2)
↓
Layer3 (6个Bottleneck, 256维, stride=2)
↓
Layer4 (3个Bottleneck, 512维, stride=2)
↓
AttentionPool2d (多头注意力)
↓
图像特征 (768维)
```
### 初始化参数
```python
layers = [3, 4, 6, 3] # ResNet50
output_dim = 768
heads = 8
image_size = 224
width = 64
```
|
|
# Bottleneck模块
### 结构组成
每个Bottleneck包含三个卷积层：
```
输入 (inplanes)
↓
1×1 Conv + BN + ReLU → planes
↓
3×3 Conv + BN + ReLU → planes
↓
AvgPool (stride>1时)
↓
1×1 Conv + BN → planes×4
↓
残差连接 (如果需要) → 输出
```
### 下采样处理
当stride>1或维度不匹配时：
```python
downsample = nn.Sequential([
nn.AvgPool2d(stride),
nn.Conv2d(inplanes, planes*4, 1, stride=1),
nn.BatchNorm2d(planes*4)
])
```
### 前向传播
```python
def forward(self, x):
identity = x
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:
identity = self.downsample(x)
out += identity # 残差连接
return self.relu3(out)
```
|
|
# 注意力池化
### AttentionPool2d
替换传统的全局平均池化，使用多头自注意力机制：
```python
class AttentionPool2d(nn.Module):
def __init__(self, spacial_dim, embed_dim, num_heads, output_dim=None):
super().__init__()
# 位置编码
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
```
### 计算流程
```
特征图 (N, C, H, W)
↓
展平为 (HW, N, C)
↓
添加CLS token (HW+1, N, C)
↓
添加位置编码
↓
多头注意力 (QKV投影)
↓
取CLS token输出 (N, output_dim)
```
### 参数
- spacial_dim：图像尺寸/32（224→7）
- embed_dim：ResNet特征维度（width×32=2048）
- num_heads：注意力头数（8）
- output_dim：输出维度（768）
|
|
# 图像预处理
### 预处理流程
```python
def image_transform(image_size=224, mean=None, std=None):
mean = mean or (0.48145466, 0.4578275, 0.40821073)
std = std or (0.26862954, 0.26130258, 0.27577711)
transforms = [
Resize(image_size, interpolation=BICUBIC),
CenterCrop(image_size),
_convert_to_rgb, # 转换为RGB
ToTensor(),
Normalize(mean=mean, std=std)
]
return Compose(transforms)
```
### RGB转换
```python
def _convert_to_rgb(image):
return image.convert('RGB')
```
### 加载图像
```python
from PIL import Image
image_path_ls = ['chest_X-ray.jpg', 'brain_MRI.jpg']
images = []
image_tensor = []
for path in image_path_ls:
image = Image.open(path).convert('RGB')
images.append(image)
image_tensor.append(preprocess_val(image))
image_tensor = torch.stack(image_tensor, dim=0).to(device)
# 输出形状: [batch_size, 3, 224, 224]
```
### 归一化参数
使用OpenAI数据集的均值和标准差：
- mean: (0.48145466, 0.4578275, 0.40821073)
- std: (0.26862954, 0.26130258, 0.27577711)
|
|
# BiomedBERT文本编码器
### 模型介绍
PMC-CLIP使用微软的BiomedBERT作为文本编码器：
- **模型**：microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract
- **用途**：医学文献摘要理解
- **优势**：在生物医学文本上预训练
### 加载模型
```python
from transformers import AutoTokenizer, AutoModel
# 加载Tokenizer
tokenizer = AutoTokenizer.from_pretrained(
'microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract'
)
# 加载模型
text_encoder = AutoModel.from_pretrained(
'microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract'
)
# 加载PMC-CLIP微调权重
text_encoder.load_state_dict(torch.load('text_encoder.pth'))
```
### Tokenization
```python
bert_input = ['chest X-ray', 'brain MRI']
encoded_input = tokenizer(
bert_input,
padding='max_length',
truncation=True,
max_length=77,
return_tensors='pt'
)
input_ids = encoded_input['input_ids'].to(device)
# 输出形状: [batch_size, max_length]
```
### 模型输出
```python
text_feature = text_encoder(input_ids)
last_hidden_state = text_feature.last_hidden_state # [batch, seq_len, 768]
pooler_output = text_feature.pooler_output # [batch, 768]
```
|
|
# 文本投影层
### 作用
将BERT的768维输出投影到与图像特征相同的共享空间：
```python
# 加载投影层权重
text_projection_layer = torch.load('text_projection_layer.pth')
text_projection_layer = nn.Parameter(text_projection_layer)
# 应用投影
text_feature = pooler_output @ text_projection_layer
# 输出: [batch_size, 768]
```
### 投影矩阵
- 输入维度：768（BERT输出）
- 输出维度：768（共享特征空间）
- 参数量：768 × 768 = 589,824
### 模型加载到设备
```python
device = 'cuda' if torch.cuda.is_available() else 'cpu'
text_encoder = text_encoder.to(device)
text_projection_layer = text_projection_layer.to(device)
```
|
|
# 文本预处理
### 分词参数
| 参数 | 值 | 说明 |
|------|-----|------|
| padding | max_length | 填充到固定长度 |
| truncation | True | 超长截断 |
| max_length | 77 | 最大序列长度 |
| return_tensors | pt | 返回PyTorch张量 |
### 医学文本示例
```python
# 医学图像描述
medical_texts = [
'chest X-ray', # 胸部X光
'brain MRI', # 脑部MRI
'CT scan of abdomen', # 腹部CT扫描
'histopathology slide' # 组织病理切片
]
```
### 注意事项
- 医学术语较多，需要专门的医学BERT
- 最大长度77足够覆盖大多数描述
- 建议添加attention_mask避免填充token干扰
|
|
# 图像特征提取
### 编码器输出
```python
image_encoder = ModifiedResNet(
layers=[3,4,6,3],
output_dim=768,
heads=8,
image_size=224,
width=64
)
# 加载预训练权重
image_encoder.load_state_dict(torch.load('image_encoder(resnet50).pth'))
image_encoder = image_encoder.to(device)
```
### 提取特征
```python
image_feature = image_encoder(image_tensor)
if isinstance(image_feature, dict):
image_feature = image_feature['image_features']
print(f'image size: {image_tensor.shape}')
print(f'feature size: {image_feature.shape}')
# 输出: image size: torch.Size([2, 3, 224, 224])
# feature size: torch.Size([2, 768])
```
### 输出说明
- 输入：[batch_size, 3, 224, 224]
- 输出：[batch_size, 768]
- 每个图像对应一个768维特征向量
|
|
# 文本特征提取
### 编码器输出
```python
text_feature = text_encoder(input_ids)
last_hidden_state = text_feature.last_hidden_state # [2, 77, 768]
pooler_output = text_feature.pooler_output # [2, 768]
```
### 应用投影层
```python
# 应用投影层
text_feature = pooler_output @ text_projection_layer
print(f'text_feature.shape: {text_feature.shape}')
# 输出: text_feature.shape: torch.Size([2, 768])
```
### 特征说明
| 输出 | 维度 | 说明 |
|------|------|------|
| last_hidden_state | [batch, seq_len, 768] | 每个token的特征 |
| pooler_output | [batch, 768] | [CLS] token的池化输出 |
| text_feature | [batch, 768] | 投影后的共享空间特征 |
|
|
# 特征归一化
### L2归一化
在计算相似度之前，需要对特征进行L2归一化：
```python
# 图像特征归一化
image_feature = image_feature / image_feature.norm(dim=-1, keepdim=True)
# 文本特征归一化
text_feature = text_feature / text_feature.norm(dim=-1, keepdim=True)
```
### 为什么需要归一化？
归一化后：
- 特征向量长度为1
- 相似度计算变为余弦相似度
- 相似度范围：[-1, 1]
- 内积 = cos(θ)
### 验证归一化
```python
# 验证归一化后的模长为1
print(f'Norm check: {image_feature[0].norm().item():.4f}')
# 应该输出: Norm check: 1.0000
```
|
|
# 相似度矩阵计算
### 计算过程
```python
# 温度参数
logit_scale = 4.4292
# 计算相似度矩阵
similarity = (math.exp(logit_scale) * (image_feature @ text_feature.T)).softmax(dim=-1)
```
### 计算步骤
1. **内积计算**：image_feature @ text_feature.T
- 输入：[2, 768] @ [768, 2]
- 输出：[2, 2] 相似度矩阵
2. **温度缩放**：math.exp(logit_scale) × 内积
- exp(4.4292) ≈ 83.87
- 放大相似度差异
3. **Softmax**：归一化为概率分布
- 每行和为1
- 表示匹配概率
### 相似度矩阵示例
```
chest X-ray brain MRI
chest_X-ray.jpg 0.98 0.02
brain_MRI.jpg 0.01 0.99
```
对角线高值表示正确的图像-文本匹配
|
|
# 温度缩放
### 温度参数
```python
logit_scale = 4.4292
temperature = math.exp(logit_scale) # ≈ 83.87
```
### 作用
温度缩放的作用是调整相似度分布的"锐度"：
- **高温度（大logit_scale）**：
- 相似度分布更陡峭
- 正确匹配的概率更高
- 训练更稳定
- **低温度（小logit_scale）**：
- 相似度分布更平缓
- 模糊度增加
### 对比
```
无温度缩放: [0.85, 0.15]
有温度缩放: [0.98, 0.02] # 更明显的区分
```
### 训练vs推理
- 训练时：温度参数可学习
- 推理时：使用固定的logit_scale
|
|
# 结果可视化
### 显示结果
```python
import matplotlib.pyplot as plt
for i, (image_path, image) in enumerate(zip(image_path_ls, images)):
print(image_path)
plt.imshow(image)
plt.show()
for j in range(len(bert_input)):
score = similarity[i, j].item()
print(f'{bert_input[j]}: {score:.4f}')
print()
```
### 输出示例
```
chest_X-ray.jpg
[显示胸部X光图像]
chest X-ray: 0.9834
brain MRI: 0.0166
brain_MRI.jpg
[显示脑部MRI图像]
chest X-ray: 0.0123
brain MRI: 0.9877
```
### 评估指标
| 指标 | 说明 |
|------|------|
| **Top-1准确率** | 最高分匹配是否正确 |
| **召回率@K** | 前K个结果中是否包含正确匹配 |
| **Mean Rank** | 正确匹配的平均排名 |
|
|
# CLIP核心原理：为什么图像和文本能对应？
## 关键问题
**图像预训练模型（ResNet）** 和 **文本预训练模型（BERT）** 是在完全不同的数据和任务上训练的，为什么在CLIP中能对应上？
## 答案：对比学习 + 联合训练
### 1. 共享特征空间
CLIP的核心思想是创建一个**共享的特征空间**：
```
训练前：图像空间 ───???─── 文本空间
(完全不同，无法比较)
训练后： 图像空间 ───→ 共享空间 ←── 文本空间
(768维对齐)
↑
对比学习训练
```
### 2. 联合训练机制
CLIP不是分别训练两个模型，而是**同时训练**：
```python
# 简化的训练伪代码
for batch in dataloader:
images, texts = batch # 配对的图像-文本数据
# 两个编码器同时前向传播
image_features = image_encoder(images) # [N, 768]
text_features = text_encoder(texts) # [N, 768]
# 计算相似度矩阵 N×N
similarity = image_features @ text_features.T
# 对比损失：让对角线（正确配对）最大
loss = contrastive_loss(similarity)
# 同时更新两个编码器
loss.backward()
optimizer.step()
```
### 3. 配对数据的作用
使用**图像-文本配对数据**进行训练：
```
训练数据示例：
┌─────────────────┬──────────────────────┐
│ 图像 │ 文本描述 │
├─────────────────┼──────────────────────┤
│ 胸部X光图像 │ "chest X-ray" │ ← 配对✓
│ 脑部MRI图像 │ "brain MRI" │ ← 配对✓
│ 胸部X光图像 │ "brain MRI" │ ← 不配对✗
└─────────────────┴──────────────────────┘
```
训练目标：配对的图像-文本特征接近，不配对的远离
|
|
# 对比学习详解
## 相似度矩阵目标
对于N个图像-文本配对，计算N×N相似度矩阵：
```
N×N 相似度矩阵示例：
正确文本 错误文本1 错误文本2
正确图像 [0.95 0.02 0.03] ← 想要最大
错误图像1 [0.01 0.96 0.03] ← 想要最小
错误图像2 [0.02 0.01 0.97] ← 想要最小
↑ ↑ ↑
拉近 推开 推开
```
## 对比损失函数
CLIP使用对称的对比损失：
```python
# 图像到文本的损失
loss_i2t = -torch.log(
torch.exp(similarity[i, i]) /
torch.exp(similarity[i]).sum()
)
# 文本到图像的损失
loss_t2i = -torch.log(
torch.exp(similarity[j, j]) /
torch.exp(similarity[:, j]).sum()
)
# 总损失
total_loss = (loss_i2t + loss_t2i) / 2
```
## 直观理解
- **拉近**：正确配对的图像和文本在特征空间中靠近
- **推远**：错误配对的图像和文本在特征空间中远离
- **结果**：模型学会将相关的图像和文本映射到相近的位置
|
|
# 投影层对齐
## 为什么需要投影层？
ResNet50的原始输出是2048维，BERT输出是768维，需要**投影层**对齐到相同维度：
```
ResNet50 输出: 2048维 ─→ projection ─→ 768维
↓
共享特征空间
↑
BERT 输出: 768维 ───→ projection ─→ 768维
```
## 投影层的作用
1. **维度对齐**：将不同维度的特征映射到统一维度
2. **语义对齐**：学习跨模态的语义映射关系
3. **可学习变换**：通过训练自动学习最优映射
## 投影层实现
```python
# 图像投影层（在AttentionPool2d中）
self.c_proj = nn.Linear(embed_dim, 768) # 2048 → 768
# 文本投影层（独立参数）
text_projection_layer = nn.Parameter(torch.randn(768, 768))
text_feature = pooler_output @ text_projection_layer # 768 → 768
```
## 投影层的训练
- 随机初始化（不使用预训练权重）
- 在CLIP训练过程中从头学习
- 学习最优的跨模态映射关系
|
|
# CLIP训练流程详解
## 完整训练流程
```
┌─────────────────────────────────────────────────┐
│ CLIP训练过程 │
└─────────────────────────────────────────────────┘
1. 准备配对数据
┌─────────────┐ ┌─────────────┐
│ Image 1 │ │ "chest X-ray" │
│ Image 2 │ │ "brain MRI" │
│ ... │ │ ... │
│ Image N │ │ "CT scan" │
└──────┬──────┘ └──────┬───────┘
2. 前向传播
│ │
▼ ▼
┌─────────┐ ┌─────────┐
│Image │ │Text │
│Encoder │ │Encoder │
│(预训练 │ │(预训练 │
│ ResNet) │ │ BERT) │
└────┬────┘ └────┬────┘
│ │
▼ ▼
┌─────────┐ ┌─────────┐
│Projection│ │Projection│
│ Layer │ │ Layer │
└────┬────┘ └────┬────┘
│ │
▼ ▼
┌────────────────────────────────┐
│ 共享特征空间 (768维) │
│ Image1: [0.2, 0.5, ..., 0.8] │
│ Text1: [0.2, 0.5, ..., 0.8] │ ← 越接近越好
│ │
│ Image2: [0.1, 0.9, ..., 0.3] │
│ Text1: [0.2, 0.5, ..., 0.8] │ ← 越远越好
└───────────────┬────────────────┘
│
▼
┌─────────────┐
│ 计算相似度 │
│ (N×N矩阵) │
└──────┬──────┘
│
▼
┌─────────────┐
│ 对比损失 │
│ 最大化对角线 │
└──────┬──────┘
│
▼
┌─────────────┐
│ 反向传播更新 │
│ 两个编码器+ │
│ 投影层 │
└─────────────┘
```
## 预训练模型的作用
| 编码器 | 预训练任务 | 对CLIP的贡献 |
|--------|-----------|-------------|
| ResNet | ImageNet分类 | 已学会提取图像特征 |
| BERT | 语言模型理解 | 已学会理解文本语义 |
**然后CLIP微调**：
- 保留预训练的知识
- 学习跨模态对齐
- 比从头训练快且效果好
|
|
# 为什么图像和文本能对应？
## 核心问题
**图像预训练模型（如ResNet）**和**文本预训练模型（如BERT）**是在完全不同的数据和任务上训练的，为什么它们还能对应上？
## 核心答案：对比学习（Contrastive Learning）
### 将图像和文本映射到共享特征空间
```
训练前：图像空间 ───???─── 文本空间
(完全不同，无法比较)
训练后： 图像空间 ───→ 共享空间 ←── 文本空间
(768维对齐)
↑
对比学习训练
```
### 为什么这不是"神奇"的？
- ✗ 不是因为模型天生理解跨模态关系
- ✓ 是通过**大量配对数据 + 对比学习训练**出来的
- ✓ 共享空间是**学习出来的**，不是预定义的
- ✓ 投影层将两个模态映射到同一维度空间
- ✓ 对比损失让正确配对接近，错误配对远离
### 共享特征空间的概念
在768维的共享空间中：
```
配对的图像-文本 → 相近的位置
↓
图像1特征: [0.2, 0.5, ..., 0.8]
文本1特征: [0.2, 0.5, ..., 0.8] ← 距离近
不配对的图像-文本 → 较远的位置
↓
图像2特征: [0.1, 0.9, ..., 0.3]
文本1特征: [0.2, 0.5, ..., 0.8] ← 距离远
```
|
|
# 联合训练机制
## CLIP不是分别训练，而是同时训练
### 训练伪代码
```python
# CLIP训练过程
for batch in dataloader:
images, texts = batch # 配对的图像-文本数据
# 两个编码器同时前向传播
image_features = image_encoder(images) # [N, 768]
text_features = text_encoder(texts) # [N, 768]
# 计算相似度矩阵 N×N
similarity = image_features @ text_features.T
# 对比损失：让对角线（正确配对）最大，非对角线最小
loss = contrastive_loss(similarity)
# 同时更新两个编码器
loss.backward()
optimizer.step()
```
### 为什么是联合训练？
| 训练方式 | 说明 | 效果 |
|---------|------|------|
| **分别训练** | 图像编码器和文本编码器独立优化 | 无法学习跨模态关系 |
| **联合训练** | 两个编码器同时优化，共享损失 | 学习到跨模态对齐 |
### 关键点
1. **同时前向传播**：图像和文本同时编码
2. **共享损失函数**：对比损失同时作用于两个模态
3. **联合反向传播**：梯度同时更新两个编码器
4. **端到端优化**：整个系统作为一个整体优化
### 与传统方法的区别
```
传统方法：图像 → 特征 → 分类器 → 类别
文本 → 特征 → 分类器 → 类别
(各自独立)
CLIP方法：图像 → 特征 → 共享空间 → 相似度 → 文本
(联合优化)
```
|
|
# 投影层对齐
## 两个预训练模型输出的维度可能不同
```
ResNet50 输出: 2048维 ─→ projection ─→ 768维
↓
共享特征空间
↑
BERT 输出: 768维 ───→ projection ─→ 768维
```
## 在代码中的实现
### 文本投影层
```python
# 加载文本投影层权重
text_projection_layer = torch.load('text_projection_layer.pth')
text_projection_layer = nn.Parameter(text_projection_layer)
# 应用投影层对齐到共享空间
pooler_output = text_encoder(input_ids).pooler_output # [N, 768]
text_feature = pooler_output @ text_projection_layer # [N, 768]
```
### 图像投影层
```python
# 在AttentionPool2d中实现
self.c_proj = nn.Linear(embed_dim, output_dim) # 2048 → 768
# 图像特征通过投影层
image_feature = self.attnpool(x) # 输出768维
```
## 投影层的三个作用
| 作用 | 说明 |
|------|------|
| **维度对齐** | 将不同维度的特征映射到统一维度 |
| **语义对齐** | 学习跨模态的语义映射关系 |
| **可学习变换** | 通过训练自动学习最优映射 |
## 投影层如何训练？
- **随机初始化**：不使用预训练权重
- **从头学习**：在CLIP训练过程中学习
- **学习目标**：找到最优的跨模态映射关系
- **结果**：将图像特征和文本特征投影到可比的空间
## 为什么投影层是关键？
没有投影层：
- ResNet输出2048维，BERT输出768维
- 无法计算相似度
- 特征空间不对齐
有了投影层：
- 统一到768维共享空间
- 可以计算相似度
- 学习跨模态语义关系
|
|
# 对比学习损失函数
## N×N 相似度矩阵示例
假设有3个图像-文本配对：
```
正确文本 错误文本1 错误文本2
正确图像 [0.95 0.02 0.03] ← 想要最大
错误图像1 [0.01 0.96 0.03] ← 想要最小
错误图像2 [0.02 0.01 0.97] ← 想要最小
↑ ↑ ↑
拉近 推开 推开
```
## 损失目标
### 对角线（正确配对）
- 目标：相似度 → 1
- 操作：拉近正样本对
### 非对角线（错误配对）
- 目标：相似度 → 0
- 操作：推开负样本对
## 对比损失实现
```python
import torch.nn.functional as F
def contrastive_loss(similarity, temperature=1.0):
"""
similarity: [N, N] 相似度矩阵
temperature: 温度参数
"""
# 温度缩放
logits = similarity / temperature
# 图像到文本的损失
labels = torch.arange(N).to(device)
loss_i2t = F.cross_entropy(logits, labels)
# 文本到图像的损失
loss_t2i = F.cross_entropy(logits.t(), labels)
# 对称损失
return (loss_i2t + loss_t2i) / 2
```
## 温度参数的作用
```python
logit_scale = 4.4292
temperature = math.exp(logit_scale) # ≈ 83.87
```
- **高温度**：相似度分布更陡峭，区分度更高
- **低温度**：相似度分布更平缓
- **可学习**：训练时作为可学习参数优化
## 直观理解
```
训练前（随机初始化）：
文本A 文本B 文本C
图像A [0.33 0.33 0.34] ← 随机分布
训练后（对比学习）：
文本A 文本B 文本C
图像A [0.95 0.02 0.03] ← 正确配对突出
```
|
|
# 训练数据
## PMC-CLIP使用的训练数据
### 数据来源：PMC（PubMed Central）
PMC是生物医学和生命科学领域的开放获取数据库。
### 数据类型：医学图像-摘要配对
| 图像类型 | 文本描述示例 |
|---------|-------------|
| 胸部X光 | "chest X-ray showing pneumonia" |
| 脑部MRI | "brain MRI scan with contrast" |
| CT扫描 | "CT scan of abdomen showing lesions" |
| 病理切片 | "histopathology of tumor tissue" |
### 配对关系
```
配对数据示例：
┌─────────────────┬──────────────────────┐
│ 图像 │ 文本描述 │
├─────────────────┼──────────────────────┤
│ 胸部X光图像 │ "chest X-ray showing │
│ │ pneumonia" │ ← 配对✓
│ │ │
│ 脑部MRI图像 │ "brain MRI scan" │ ← 配对✓
│ │ │
│ 胸部X光图像 │ "brain MRI scan" │ ← 不配对✗
└─────────────────┴──────────────────────┘
```
### 训练数据的作用
1. **提供监督信号**：图像-文本配对告诉模型什么是正确的对应
2. **学习跨模态关系**：模型从大量配对中学习图像和文本的关联
3. **医学领域知识**：PMC数据确保模型理解医学图像和术语
### 数据规模
- CLIP原版：4亿图像-文本对
- PMC-CLIP：百万级医学图像-摘要对
- 数据质量：经过同行评审的医学文献
### 为什么配对数据重要？
没有配对数据：
- 模型不知道图像和文本应该如何对应
- 无法计算对比损失
- 无法学习跨模态对齐
有配对数据：
- 提供明确的监督信号
- 可以计算对比损失
- 可以学习跨模态语义关系
|
|
# 完整训练流程图
## CLIP训练过程详解
```
┌─────────────────────────────────────────────────────┐
│ CLIP训练过程 │
└─────────────────────────────────────────────────────┘
第1步：准备配对数据
┌─────────────┐ ┌─────────────┐
│ Image 1 │ │ "chest X-ray"│
│ Image 2 │ │ "brain MRI" │
│ ... │ │ ... │
│ Image N │ │ "CT scan" │
└──────┬──────┘ └──────┬───────┘
│ │
第2步：编码器前向传播
▼ ▼
┌─────────────┐ ┌─────────────┐
│ Image │ │ Text │
│ Encoder │ │ Encoder │
│ (预训练ResNet)│ │ (预训练BERT) │
└──────┬──────┘ └──────┬───────┘
│ │
▼ ▼
┌─────────────┐ ┌─────────────┐
│ Projection │ │ Projection │
│ Layer │ │ Layer │
│ (2048→768) │ │ (768→768) │
└──────┬──────┘ └──────┬───────┘
│ │
第3步：共享特征空间
▼ ▼
┌───────────────────────────────────────┐
│ 共享特征空间 (768维) │
│ Image1: [0.2, 0.5, ..., 0.8] │
│ Text1: [0.2, 0.5, ..., 0.8] ← 越接近越好│
│ │
│ Image2: [0.1, 0.9, ..., 0.3] │
│ Text1: [0.2, 0.5, ..., 0.8] ← 越远越好 │
└───────────────────┬─────────────────┘
│
第4步：计算相似度
▼
┌─────────────┐
│ 计算相似度 │
│ (N×N矩阵) │
└──────┬──────┘
│
第5步：计算对比损失
▼
┌─────────────┐
│ 对比损失 │
│ 最大化对角线 │
│ 最小化非对角 │
└──────┬──────┘
│
第6步：反向传播更新
▼
┌─────────────┐
│ 反向传播更新 │
│ 两个编码器 │
│ + 投影层 │
└─────────────┘
```
## 各步骤详解
### 步骤1：数据准备
- 收集图像-文本配对数据
- 确保配对质量
- 构建训练批次
### 步骤2：编码器前向传播
- 图像通过ResNet提取特征
- 文本通过BERT提取特征
- 保持预训练知识
### 步骤3：投影到共享空间
- 图像特征通过投影层（2048→768）
- 文本特征通过投影层（768→768）
- 统一到可比的维度
### 步骤4：计算相似度
- 计算N×N相似度矩阵
- 每个图像与所有文本比较
- 每个文本与所有图像比较
### 步骤5：计算对比损失
- 对角线（正确配对）最大化
- 非对角线（错误配对）最小化
- 对称计算（图像→文本，文本→图像）
### 步骤6：反向传播更新
- 同时更新两个编码器
- 同时更新投影层
- 保持预训练知识的同时学习跨模态对齐
|
|
# 为什么预训练模型还能用？
## 预训练模型提供了良好的初始化
### 预训练模型对CLIP的贡献
| 编码器 | 预训练任务 | 对CLIP的贡献 |
|--------|-----------|-------------|
| **ResNet** | ImageNet图像分类 | 已学会提取图像特征（边缘→纹理→物体） |
| **BERT** | 语言模型理解 | 已学会理解文本语义 |
### 预训练学到什么？
#### ResNet学到的图像特征
```
浅层：边缘、颜色、纹理
中层：局部模式、部件
深层：整体物体、场景
```
#### BERT学到的文本特征
```
浅层：语法、词性
中层：短语、句法
深层：语义、上下文理解
```
## CLIP的微调过程
### 训练策略
1. **加载预训练权重**：初始化编码器
2. **随机初始化投影层**：从头学习跨模态映射
3. **联合训练**：同时优化两个编码器
4. **学习率调整**：编码器用较小学习率，投影层用较大学习率
### 微调效果
```
预训练模型（如ResNet）
↓
已具备特征提取能力
↓
CLIP微调
↓
保持特征提取能力 + 学习跨模态对齐
↓
比从头训练快得多、效果好得多
```
## 为什么不用从头训练？
| 对比 | 从头训练 | 基于预训练微调 |
|------|---------|--------------|
| **训练数据** | 需要更多数据 | 需要较少数据 |
| **训练时间** | 更长 | 更短 |
| **收敛速度** | 更慢 | 更快 |
| **最终效果** | 可能更差 | 通常更好 |
## 预训练+微调的优势
1. **特征提取能力**：预训练模型已经学会提取高质量特征
2. **知识迁移**：将大规模预训练的知识迁移到CLIP任务
3. **快速收敛**：好的初始化使训练更快收敛
4. **更好效果**：预训练的知识提供更强的泛化能力
## 总结
**不是神奇**：是大量配对数据 + 对比学习训练出来的
**共享空间**：投影层将两个模态映射到同一维度空间
**对比损失**：让正确配对接近，错误配对远离
**预训练基础**：利用了ResNet和BERT已有的特征提取能力
**联合训练**：同时优化两个编码器学习跨模态对齐
|
|
# 环境配置
## 导入依赖库
```python
# 基础库
from collections import OrderedDict
from typing import Optional, Sequence, Tuple
import math
# PyTorch相关
import torch
import torch.nn.functional as F
from torch import nn
# 图像处理
from torchvision.transforms import (
Normalize, Compose, RandomResizedCrop,
InterpolationMode, ToTensor, Resize, CenterCrop
)
# Transformers库
from transformers import AutoTokenizer, AutoModel
# 图像显示
from PIL import Image
import matplotlib.pyplot as plt
```
## 检查环境
```bash
# 检查PyTorch版本
!pip show torch
# Name: torch
# Version: 2.6.0+cu124
# 检查Transformers版本
!pip show transformers
# Name: transformers
# Version: 4.54.0
```
## 环境要求
| 组件 | 版本要求 | 说明 |
|------|---------|------|
| Python | >= 3.8 | 推荐使用3.11 |
| PyTorch | >= 2.0.0 | 支持CUDA加速 |
| Transformers | >= 4.54.0 | 用于加载BERT |
| CUDA | 推荐有 | GPU加速推理 |
## 设备配置
```python
# 自动检测设备
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'Using device: {device}')
```
|
|
# ResNet架构实现
## Bottleneck模块
```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时的平均池化
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)
# 下采样（残差连接）
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: torch.Tensor):
identity = x
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
```
## 注意力池化层
```python
class AttentionPool2d(nn.Module):
def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
super().__init__()
# 位置编码
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.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3]).permute(2, 0, 1)
# 添加CLS token
x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)
# 添加位置编码
x = x + self.positional_embedding[:, None, :].to(x.dtype)
# 多头注意力
x, _ = F.multi_head_attention_forward(
query=x, 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]),
bias_k=None, bias_v=None, add_zero_attn=False,
dropout_p=0, 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[0] # 返回CLS token
```
## ModifiedResNet完整实现
```python
class ModifiedResNet(nn.Module):
def __init__(self, layers, output_dim, heads, image_size=224, width=64):
super().__init__()
self.output_dim = output_dim
self.image_size = image_size
# 3层Stem
self.conv1 = nn.Conv2d(3, width // 2, kernel_size=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, kernel_size=3, stride=2, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(width // 2)
self.relu2 = nn.ReLU(inplace=True)
self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, stride=2, padding=1, bias=False)
self.bn3 = nn.BatchNorm2d(width)
self.relu3 = nn.ReLU(inplace=True)
self.avgpool = nn.AvgPool2d(2)
# 残差层
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)
# 注意力池化
embed_dim = width * 32 # ResNet特征维度
self.attnpool = AttentionPool2d(image_size // 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):
x = self.stem(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.attnpool(x)
visual_output = dict.fromkeys(["image_features", "mim_loss"], None)
visual_output.update({
'image_features': x,
})
return visual_output
def stem(self, 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
```
|
|
# 图像预处理
## RGB转换函数
```python
def _convert_to_rgb(image):
"""将图像转换为RGB格式"""
return image.convert('RGB')
```
## 图像变换流程
```python
def image_transform(
image_size: int,
mean: Optional[Tuple[float, ...]] = None,
std: Optional[Tuple[float, ...]] = None,
fill_color: int = 0,
):
"""
图像预处理变换
Args:
image_size: 目标图像尺寸
mean: 归一化均值
std: 归一化标准差
fill_color: 填充颜色
"""
# 处理方形尺寸
if isinstance(image_size, (list, tuple)) and image_size[0] == image_size[1]:
image_size = image_size[0]
# 使用OpenAI数据集的统计值
mean = mean or (0.48145466, 0.4578275, 0.40821073)
std = std or (0.26862954, 0.26130258, 0.27577711)
normalize = Normalize(mean=mean, std=std)
# 构建变换序列
transforms = [
Resize(image_size, interpolation=InterpolationMode.BICUBIC),
CenterCrop(image_size),
_convert_to_rgb,
ToTensor(),
normalize,
]
return Compose(transforms)
```
## 加载图像
```python
from PIL import Image
# 定义图像路径
image_path_ls = [
'chest_X-ray.jpg',
'brain_MRI.jpg'
]
# 创建预处理器
preprocess_val = image_transform(image_size=224)
# 加载图像
images = []
image_tensor = []
for image_path in image_path_ls:
image = Image.open(image_path).convert('RGB')
images.append(image)
image_tensor.append(preprocess_val(image))
# 堆叠为batch
image_tensor = torch.stack(image_tensor, dim=0).to(device)
print(f'Image tensor shape: {image_tensor.shape}')
# 输出: Image tensor shape: torch.Size([2, 3, 224, 224])
```
## 下载测试图像
```bash
# 下载胸部X光图像
!wget https://huggingface.co/datasets/axiong/pmc_oa/resolve/main/chest_X-ray.jpg
# 下载脑部MRI图像
!wget https://huggingface.co/datasets/axiong/pmc_oa/resolve/main/brain_MRI.jpg
```
|
|
# 加载模型
## 下载模型权重
```bash
# 下载文本编码器（约418MB）
!wget https://huggingface.co/datasets/axiong/pmc_oa/resolve/main/text_encoder.pth
# 下载图像编码器 ResNet50（约144MB）
!wget https://huggingface.co/datasets/axiong/pmc_oa/resolve/main/image_encoder(resnet50).pth
# 下载文本投影层（约2.3MB）
!wget https://huggingface.co/datasets/axiong/pmc_oa/resolve/main/text_projection_layer.pth
```
## 初始化图像编码器
```python
# 创建ModifiedResNet模型
image_encoder = ModifiedResNet(
layers=[3, 4, 6, 3], # ResNet50结构
output_dim=768, # 输出维度
heads=8, # 注意力头数
image_size=224, # 输入图像尺寸
width=64 # 基础宽度
)
# 加载预训练权重
image_encoder.load_state_dict(torch.load('image_encoder(resnet50).pth'))
# 移到设备
image_encoder = image_encoder.to(device)
```
## 初始化文本编码器
```python
# 加载BiomedBERT分词器
tokenizer = AutoTokenizer.from_pretrained(
'microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract'
)
# 加载BiomedBERT模型
text_encoder = AutoModel.from_pretrained(
'microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract'
)
# 加载PMC-CLIP微调权重
text_encoder.load_state_dict(torch.load('text_encoder.pth'))
# 移到设备
text_encoder = text_encoder.to(device)
```
## 加载投影层和温度参数
```python
# 加载文本投影层
text_projection_layer = torch.load('text_projection_layer.pth')
text_projection_layer = nn.Parameter(text_projection_layer)
text_projection_layer = text_projection_layer.to(device)
# 设置温度参数
logit_scale = 4.4292
```
## 模型参数汇总
| 组件 | 参数量 | 文件大小 |
|------|-------|---------|
| 图像编码器 (ResNet50) | ~25M | 144MB |
| 文本编码器 (BiomedBERT) | ~110M | 418MB |
| 文本投影层 | 0.59M | 2.3MB |
| **总计** | **~135M** | **~564MB** |
|
|
# 特征提取
## 图像特征提取
```python
# 提取图像特征
image_feature = image_encoder(image_tensor)
# 处理字典输出
if isinstance(image_feature, dict):
image_feature = image_feature['image_features']
print(f'Image size: {image_tensor.shape}')
print(f'Feature size: {image_feature.shape}')
# 输出: Image size: torch.Size([2, 3, 224, 224])
# Feature size: torch.Size([2, 768])
```
## 文本特征提取
```python
# 准备文本输入
bert_input = [
'chest X-ray',
'brain MRI',
]
# 分词
encoded_input = tokenizer(
bert_input,
padding='max_length',
truncation=True,
max_length=77,
return_tensors='pt'
)
input_ids = encoded_input['input_ids'].to(device)
# 文本编码
text_feature = text_encoder(input_ids)
last_hidden_state = text_feature.last_hidden_state # [2, 77, 768]
pooler_output = text_feature.pooler_output # [2, 768]
# 应用投影层
text_feature = pooler_output @ text_projection_layer
print(f'Text feature shape: {text_feature.shape}')
# 输出: Text feature shape: torch.Size([2, 768])
```
## 特征归一化
```python
# L2归一化（重要！）
image_feature = image_feature / image_feature.norm(dim=-1, keepdim=True)
text_feature = text_feature / text_feature.norm(dim=-1, keepdim=True)
# 验证归一化
print(f'Image norm: {image_feature[0].norm().item():.4f}')
print(f'Text norm: {text_feature[0].norm().item():.4f}')
# 输出: Image norm: 1.0000
# Text norm: 1.0000
```
## 特征提取流程图
```
图像输入 [2, 3, 224, 224]
↓
ResNet编码器
↓
图像特征 [2, 768]
↓
L2归一化
↓
归一化图像特征 [2, 768]
文本输入 ["chest X-ray", "brain MRI"]
↓
Tokenizer
↓
input_ids [2, 77]
↓
BERT编码器
↓
pooler_output [2, 768]
↓
投影层
↓
文本特征 [2, 768]
↓
L2归一化
↓
归一化文本特征 [2, 768]
```
|
|
# 相似度计算与可视化
## 计算相似度矩阵
```python
import math
# 温度缩放并计算相似度
similarity = (math.exp(logit_scale) * image_feature @ text_feature.T).softmax(dim=-1)
print(f'Similarity matrix shape: {similarity.shape}')
print(f'Temperature: {math.exp(logit_scale):.2f}')
# 输出: Similarity matrix shape: torch.Size([2, 2])
# Temperature: 83.87
```
## 显示相似度矩阵
```python
# 打印相似度矩阵
print("\n相似度矩阵:")
print(similarity.detach().cpu().numpy())
# 示例输出:
# [[0.9834 0.0166]
# [0.0123 0.9877]]
```
## 结果可视化
```python
import matplotlib.pyplot as plt
# 遍历每个图像
for i, (image_path, image) in enumerate(zip(image_path_ls, images)):
print(f'\n{'='*50}')
print(f'图像: {image_path}')
print(f'{'='*50}')
# 显示图像
plt.figure(figsize=(6, 6))
plt.imshow(image)
plt.axis('off')
plt.title(image_path)
plt.show()
# 打印相似度分数
print(f'\n相似度分数:')
for j in range(len(bert_input)):
score = similarity[i, j].item()
# 添加颜色标记
if i == j: # 正确配对
print(f' ✓ {bert_input[j]}: {score:.4f}')
else: # 错误配对
print(f' ✗ {bert_input[j]}: {score:.4f}')
```
## 输出示例
```
==================================================
图像: chest_X-ray.jpg
==================================================
[显示胸部X光图像]
相似度分数:
✓ chest X-ray: 0.9834
✗ brain MRI: 0.0166
==================================================
图像: brain_MRI.jpg
==================================================
[显示脑部MRI图像]
相似度分数:
✗ chest X-ray: 0.0123
✓ brain MRI: 0.9877
```
## 完整推理代码
```python
def compute_similarity(image_paths, text_inputs):
"""
计算图像-文本相似度
Args:
image_paths: 图像路径列表
text_inputs: 文本描述列表
Returns:
similarity: 相似度矩阵
"""
# 加载图像
images = []
image_tensors = []
for path in image_paths:
img = Image.open(path).convert('RGB')
images.append(img)
image_tensors.append(preprocess_val(img))
image_tensor = torch.stack(image_tensors, dim=0).to(device)
# 提取特征
with torch.no_grad():
image_feat = image_encoder(image_tensor)
if isinstance(image_feat, dict):
image_feat = image_feat['image_features']
# 文本编码
encoded = tokenizer(text_inputs, padding='max_length',
truncation=True, max_length=77, return_tensors='pt')
input_ids = encoded['input_ids'].to(device)
text_feat = text_encoder(input_ids).pooler_output
text_feat = text_feat @ text_projection_layer
# 归一化
image_feat = image_feat / image_feat.norm(dim=-1, keepdim=True)
text_feat = text_feat / text_feat.norm(dim=-1, keepdim=True)
# 计算相似度
similarity = (math.exp(logit_scale) * image_feat @ text_feat.T).softmax(dim=-1)
return similarity, images
```
|
|
# OpenAI CLIP 概述
## 背景
2021年1月，OpenAI宣布了两个多模态模型：**DALL-E**和**CLIP**，两者都以某种方式连接文本和图像。本文我们将从头实现CLIP模型。
## 什么是CLIP？
**CLIP**全称**Contrastive Language-Image Pre-training**（对比语言-图像预训练）。
### 核心功能
CLIP学习**整句句子**与其描述的**图像**之间的关系。训练完成后：
- 输入一个句子，模型能够检索最相关的图像
- 输入一张图像，模型能够找到最相关的描述
### 与传统方法的区别
| 传统图像分类 | CLIP方法 |
|-------------|----------|
| 单个类别标签（car, dog等） | 完整的自然语言描述 |
| 在ImageNet上训练优化分类 | 在图像-文本对上学习 |
| 需要重新训练才能识别新类别 | 零样本即可识别新类别 |
## 为什么CLIP有趣？
### 零样本能力
CLIP在大规模数据集上训练后，可以无需额外训练就作为分类器使用。甚至比专门在ImageNet上优化的SOTA模型表现更好！
### 学习更丰富的表示
在完整句子上训练，模型可以学到更多东西，发现图像和文本之间的模式。
## 示例效果
输入查询文本：
- "a boy jumping with skateboard"
- "a girl jumping from swing"
模型可以检索出最相关的图像。
## 论文信息
**标题**: Learning Transferable Visual Models From Natural Language Supervision
**作者**: OpenAI (Alec Radford et al.)
**年份**: 2021
**亮点**: 大规模对比学习、零样本能力、跨模态迁移
|
|
# CLIP核心思想
## 学习目标
**我们希望模型学习：为给定图像和描述它的文本，生成"相似的表示（向量）"。**
也就是说，无论我们给它图像还是描述它的文本，我们都希望它为两者产生相同大小的向量。
## 直观理解
```
输入: 图像A + 描述文本T
↓
图像编码器 → 图像特征向量 [256维]
文本编码器 → 文本特征向量 [256维]
↓
两个向量应该非常相似（接近）
```
## 如何判断两个向量相似？
在线性代数中，衡量两个向量是否具有相似特征的常用方法是计算它们的**点积**（内积）：
- 如果点积值大 → 它们相似
- 如果点积值小 → 它们不相似
## CLIP的工作原理
### 1. 编码阶段
```
图像 → 图像编码器 → 2048维特征 → 投影层 → 256维嵌入
文本 → 文本编码器 → 768维特征 → 投影层 → 256维嵌入
```
### 2. 对比学习
- **拉近**: 匹配的图像-文本对的嵌入应该接近
- **推远**: 不匹配的图像-文本对的嵌入应该远离
### 3. 相似度计算
```
相似度矩阵 = text_embeddings @ image_embeddings.T
= (batch_size, 256) @ (256, batch_size)
= (batch_size, batch_size)
```
## 训练策略
### 数据规模
- CLIP原版在**4亿**图像-文本对上训练
- 使用来自互联网的自然图像和文本
### 为什么大规模重要？
- 更多样化的视觉概念
- 更丰富的语言描述
- 更好的泛化能力
## 理想情况下的相似度矩阵
如果模型训练完美，相似度矩阵的softmax结果应该是**单位矩阵**：
```
[[1.0, 0.0, 0.0, 0.0],
[0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 1.0, 0.0],
[0.0, 0.0, 0.0, 1.0]]
```
对角线为1.0（正确匹配），其余为0.0（错误匹配）。
|
|
# CLIP模型架构
## 整体架构
```
┌─────────────────────────────────────────────────┐
│ CLIP模型架构 │
└─────────────────────────────────────────────────┘
图像输入 文本输入
(224×224×3) ("a dog on grass")
↓ ↓
┌───────────┐ ┌───────────┐
│Image │ │Text │
│Encoder │ │Encoder │
│(ResNet50) │ │(DistilBERT)│
└─────┬─────┘ └─────┬─────┘
│ │
↓ ↓
图像特征 文本特征
[batch, 2048] [batch, 768]
│ │
↓ ↓
┌───────────┐ ┌───────────┐
│Image │ │Text │
│Projection │ │Projection │
│Head │ │Head │
└─────┬─────┘ └─────┬─────┘
│ │
↓ ↓
图像嵌入 文本嵌入
[batch, 256] [batch, 256]
│ │
└──────┬───────────────────┘
↓
计算相似度矩阵
[batch, batch]
↓
对比损失
```
## 图像编码器
### 使用timm库创建ResNet50
```python
import timm
import torch.nn as nn
class ImageEncoder(nn.Module):
def __init__(self, model_name='resnet50', pretrained=True, trainable=True):
super().__init__()
self.model = timm.create_model(
model_name,
pretrained,
num_classes=0, # 移除分类头
global_pool="avg" # 全局平均池化
)
for p in self.model.parameters():
p.requires_grad = trainable
def forward(self, x):
return self.model(x) # 输出: [batch, 2048]
```
### 图像编码器特点
| 特性 | 说明 |
|------|------|
| **模型** | ResNet50 |
| **预训练** | ImageNet |
| **输出维度** | 2048 |
| **可训练** | 可配置 |
## 文本编码器
### 使用DistilBERT
```python
from transformers import DistilBertModel
class TextEncoder(nn.Module):
def __init__(self, model_name="distilbert-base-uncased",
pretrained=True, trainable=True):
super().__init__()
if pretrained:
self.model = DistilBertModel.from_pretrained(model_name)
else:
self.model = DistilBertModel(config=DistilBertConfig())
for p in self.model.parameters():
p.requires_grad = trainable
# 使用CLS token的表示作为句子嵌入
self.target_token_idx = 0
def forward(self, input_ids, attention_mask):
output = self.model(input_ids=input_ids,
attention_mask=attention_mask)
last_hidden_state = output.last_hidden_state
return last_hidden_state[:, self.target_token_idx, :]
# 输出: [batch, 768]
```
### 文本编码器特点
| 特性 | 说明 |
|------|------|
| **模型** | DistilBERT |
| **输出维度** | 768 |
| **使用Token** | CLS token |
| **原因** | CLS token捕捉整个句子的表示 |
## 投影头
### 为什么要投影？
图像编码器输出2048维，文本编码器输出768维，维度不同无法直接比较。投影头将它们映射到相同的256维空间。
### 投影头实现
```python
class ProjectionHead(nn.Module):
def __init__(self, embedding_dim, projection_dim=256, dropout=0.1):
super().__init__()
self.projection = nn.Linear(embedding_dim, projection_dim)
self.gelu = nn.GELU()
self.fc = nn.Linear(projection_dim, projection_dim)
self.dropout = nn.Dropout(dropout)
self.layer_norm = nn.LayerNorm(projection_dim)
def forward(self, x):
projected = self.projection(x)
x = self.gelu(projected)
x = self.fc(x)
x = self.dropout(x)
x = x + projected # 残差连接
x = self.layer_norm(x)
return x # 输出: [batch, 256]
```
### 投影头结构
```
输入 [embedding_dim]
↓
Linear → [projection_dim]
↓
GELU激活
↓
Linear → [projection_dim]
↓
Dropout
↓
残差连接 (x + projected)
↓
LayerNorm
↓
输出 [projection_dim]
```
## 完整CLIP模型
```python
class CLIPModel(nn.Module):
def __init__(self, temperature=1.0,
image_embedding=2048,
text_embedding=768):
super().__init__()
self.image_encoder = ImageEncoder()
self.text_encoder = TextEncoder()
self.image_projection = ProjectionHead(embedding_dim=image_embedding)
self.text_projection = ProjectionHead(embedding_dim=text_embedding)
self.temperature = temperature
def forward(self, batch):
# 获取图像和文本特征
image_features = self.image_encoder(batch["image"])
text_features = self.text_encoder(
input_ids=batch["input_ids"],
attention_mask=batch["attention_mask"]
)
# 获取图像和文本嵌入（相同维度）
image_embeddings = self.image_projection(image_features)
text_embeddings = self.text_projection(text_features)
# 计算损失
logits = (text_embeddings @ image_embeddings.T) / self.temperature
# 计算相似度矩阵作为目标
images_similarity = image_embeddings @ image_embeddings.T
texts_similarity = text_embeddings @ text_embeddings.T
targets = F.softmax(
(images_similarity + texts_similarity) / 2 * self.temperature,
dim=-1
)
# 计算交叉熵损失
texts_loss = cross_entropy(logits, targets, reduction='none')
images_loss = cross_entropy(logits.T, targets.T, reduction='none')
loss = (images_loss + texts_loss) / 2.0
return loss.mean()
```
|
|
# CLIP对比损失
## 损失函数原理
### 目标矩阵
理想情况下，相似度矩阵的softmax应该是单位矩阵：
```python
# 示例：batch_size=4, dim=256
batch_size = 4
dim = 256
embeddings = torch.randn(batch_size, dim)
similarity_matrix = embeddings @ embeddings.T
# softmax后
print(F.softmax(similarity_matrix, dim=-1))
# 输出接近单位矩阵
```
### 为什么这样设计？
- **对角线元素=1**: 正确配对的图像-文本完全匹配
- **非对角线元素=0**: 错误配对完全不匹配
## 损失计算步骤
### 1. 计算相似度矩阵
```python
# logits形状: [batch_size, batch_size]
logits = (text_embeddings @ image_embeddings.T) / temperature
```
### 2. 计算目标矩阵
```python
# 图像之间的相似度
images_similarity = image_embeddings @ image_embeddings.T
# 文本之间的相似度
texts_similarity = text_embeddings @ text_embeddings.T
# 合并作为目标
targets = F.softmax(
(images_similarity + texts_similarity) / 2 * temperature,
dim=-1
)
```
### 3. 交叉熵损失
```python
def cross_entropy(preds, targets, reduction='none'):
log_softmax = nn.LogSoftmax(dim=-1)
loss = (-targets * log_softmax(preds)).sum(1)
if reduction == "none":
return loss
elif reduction == "mean":
return loss.mean()
# 文本到图像的损失
texts_loss = cross_entropy(logits, targets, reduction='none')
# 图像到文本的损失
images_loss = cross_entropy(logits.T, targets.T, reduction='none')
# 对称损失
loss = (images_loss + texts_loss) / 2.0
```
## 温度参数
### 温度的作用
```python
temperature = 1.0 # 可学习参数
logits = (text_embeddings @ image_embeddings.T) / temperature
```
| 温度值 | 效果 |
|--------|------|
| **高温度** | 相似度分布更平缓，训练更稳定 |
| **低温度** | 相似度分布更陡峭，区分度更高 |
| **temperature=1.0** | 无缩放效果 |
## 为什么不简化为torch.arange？
更简单的做法是：
```python
# 简单版本（不适用于多标签情况）
nn.CrossEntropyLoss()(logits, torch.arange(batch_size))
```
### 为什么不使用简化版本？
1. **多标签数据集**: 一张图像可能有多条描述
2. **批内重复**: 同一图像可能多次出现在同一batch
3. **正确处理**: 完整目标矩阵可以正确处理这些情况
## 损失函数可视化
```
图像1 图像2 图像3
文本1 [0.95 0.03 0.02] ← 高（正确）
文本2 [0.02 0.96 0.02] ← 高（正确）
文本3 [0.03 0.01 0.96] ← 高（正确）
```
训练目标：让对角线接近1，非对角线接近0。
|
|
# CLIP训练实现
## 配置参数
```python
class CFG:
# 训练参数
batch_size = 32
epochs = 4
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 学习率
head_lr = 1e-3
image_encoder_lr = 1e-4
text_encoder_lr = 1e-5
weight_decay = 1e-3
# 模型参数
model_name = 'resnet50'
image_embedding = 2048
text_encoder_model = "distilbert-base-uncased"
text_embedding = 768
max_length = 200
# 投影头参数
num_projection_layers = 1
projection_dim = 256
dropout = 0.1
temperature = 1.0
# 图像大小
size = 224
```
## 数据集实现
```python
class CLIPDataset(torch.utils.data.Dataset):
def __init__(self, image_filenames, captions, tokenizer, transforms):
self.image_filenames = image_filenames
self.captions = list(captions)
self.encoded_captions = tokenizer(
list(captions),
padding=True,
truncation=True,
max_length=CFG.max_length
)
self.transforms = transforms
def __getitem__(self, idx):
item = {
key: torch.tensor(values[idx])
for key, values in self.encoded_captions.items()
}
# 加载并处理图像
image = cv2.imread(f"{CFG.image_path}/{self.image_filenames[idx]}")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = self.transforms(image=image)['image']
item['image'] = torch.tensor(image).permute(2, 0, 1).float()
item['caption'] = self.captions[idx]
return item
def __len__(self):
return len(self.captions)
```
## 图像增强
```python
import albumentations as A
def get_transforms(mode="train"):
return A.Compose([
A.Resize(CFG.size, CFG.size, always_apply=True),
A.Normalize(max_pixel_value=255.0, always_apply=True),
])
```
## 训练循环
```python
def train_epoch(model, train_loader, optimizer, lr_scheduler, step):
loss_meter = AvgMeter()
tqdm_object = tqdm(train_loader, total=len(train_loader))
for batch in tqdm_object:
batch = {k: v.to(CFG.device)
for k, v in batch.items() if k != "caption"}
# 前向传播
loss = model(batch)
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
if step == "batch":
lr_scheduler.step()
# 更新损失
count = batch["image"].size(0)
loss_meter.update(loss.item(), count)
tqdm_object.set_postfix(
train_loss=loss_meter.avg,
lr=get_lr(optimizer)
)
return loss_meter
```
## 优化器配置
```python
import itertools
def main():
# 创建数据加载器
train_df, valid_df = make_train_valid_dfs()
tokenizer = DistilBertTokenizer.from_pretrained(CFG.text_tokenizer)
train_loader = build_loaders(train_df, tokenizer, mode="train")
valid_loader = build_loaders(valid_df, tokenizer, mode="valid")
# 创建模型
model = CLIPModel().to(CFG.device)
# 配置不同层的学习率
params = [
{
"params": model.image_encoder.parameters(),
"lr": CFG.image_encoder_lr
},
{
"params": model.text_encoder.parameters(),
"lr": CFG.text_encoder_lr
},
{
"params": itertools.chain(
model.image_projection.parameters(),
model.text_projection.parameters()
),
"lr": CFG.head_lr,
"weight_decay": CFG.weight_decay
}
]
# 优化器和调度器
optimizer = torch.optim.AdamW(params, weight_decay=0.)
lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer, mode="min",
patience=1, factor=0.8
)
# 训练循环
best_loss = float('inf')
for epoch in range(CFG.epochs):
print(f"Epoch: {epoch + 1}")
model.train()
train_loss = train_epoch(model, train_loader, optimizer, lr_scheduler, "epoch")
model.eval()
with torch.no_grad():
valid_loss = valid_epoch(model, valid_loader)
# 保存最佳模型
if valid_loss.avg < best_loss:
best_loss = valid_loss.avg
torch.save(model.state_dict(), "best.pt")
print("Saved Best Model!")
lr_scheduler.step(valid_loss.avg)
```
## 训练技巧
| 技巧 | 说明 |
|------|------|
| **不同学习率** | 编码器用小学习率，投影头用大学习率 |
| **权重衰减** | 防止过拟合 |
| **学习率调度** | ReduceLROnPlateau自适应调整 |
| **早停** | 保存验证损失最低的模型 |
|
|
# CLIP推理应用
## 获取图像嵌入
```python
def get_image_embeddings(valid_df, model_path):
tokenizer = DistilBertTokenizer.from_pretrained(CFG.text_tokenizer)
valid_loader = build_loaders(valid_df, tokenizer, mode="valid")
model = CLIPModel().to(CFG.device)
model.load_state_dict(torch.load(model_path, map_location=CFG.device))
model.eval()
valid_image_embeddings = []
with torch.no_grad():
for batch in tqdm(valid_loader):
image_features = model.image_encoder(batch["image"].to(CFG.device))
image_embeddings = model.image_projection(image_features)
valid_image_embeddings.append(image_embeddings)
return model, torch.cat(valid_image_embeddings)
```
## 查找匹配图像
```python
def find_matches(model, image_embeddings, query,
image_filenames, n=9):
tokenizer = DistilBertTokenizer.from_pretrained(CFG.text_tokenizer)
# 编码查询文本
encoded_query = tokenizer([query])
batch = {
key: torch.tensor(values).to(CFG.device)
for key, values in encoded_query.items()
}
with torch.no_grad():
text_features = model.text_encoder(
input_ids=batch["input_ids"],
attention_mask=batch["attention_mask"]
)
text_embeddings = model.text_projection(text_features)
# 归一化
image_embeddings_n = F.normalize(image_embeddings, p=2, dim=-1)
text_embeddings_n = F.normalize(text_embeddings, p=2, dim=-1)
# 计算相似度
dot_similarity = text_embeddings_n @ image_embeddings_n.T
# 获取top-k
values, indices = torch.topk(dot_similarity.squeeze(0), n * 5)
# 显示结果
matches = [image_filenames[idx] for idx in indices[::5]]
_, axes = plt.subplots(3, 3, figsize=(10, 10))
for match, ax in zip(matches, axes.flatten()):
image = cv2.imread(f"{CFG.image_path}/{match}")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
ax.imshow(image)
ax.axis("off")
plt.show()
```
## 使用示例
```python
# 获取验证集图像嵌入
_, valid_df = make_train_valid_dfs()
model, image_embeddings = get_image_embeddings(valid_df, "best.pt")
# 查询
find_matches(
model,
image_embeddings,
query="dogs on the grass",
image_filenames=valid_df['image'].values,
n=9
)
```
## 应用场景
### 1. 图像检索
输入文本描述，检索相关图像：
- "dogs on the grass" → 找到草地上的狗的图片
- "a boy jumping with skateboard" → 找到男孩玩滑板的图片
### 2. 零样本分类
```python
# 定义类别文本提示
class_prompts = [
"a photo of a dog",
"a photo of a cat",
"a photo of a bird",
"a photo of a car"
]
# 编码类别文本
text_features = encode_text(class_prompts)
# 编码测试图像
image_features = encode_image(test_image)
# 找到最匹配的类别
similarity = image_features @ text_features.T
predicted_class = class_prompts[similarity.argmax()]
```
### 3. 图像-文本匹配
验证图像和描述是否匹配：
- 图像检索系统
- 自动标注
- 内容审核
## 性能优化
| 技术 | 说明 |
|------|------|
| **批处理** | 一次处理多个查询 |
| **归一化** | 使用余弦相似度提高稳定性 |
| **缓存嵌入** | 预先计算并存储图像嵌入 |
| **近似搜索** | 使用FAISS等加速大规模检索 |
## CLIP原版vs实现版本
| 特性 | OpenAI CLIP | 实现版本 |
|------|-------------|----------|
| **图像编码器** | Vision Transformer | ResNet50 |
| **文本编码器** | Custom Transformer | DistilBERT |
| **训练数据** | 4亿对 | Flickr 8k/30k |
| **投影维度** | 512 | 256 |
| **核心思想** | 相同 | 相同 |
|
|
# 环境配置与导入
## 安装依赖
```bash
pip install timm
pip install transformers
```
## 导入库
```python
# 基础库
import os
import gc
import cv2
import numpy as np
import pandas as pd
import itertools
from tqdm.autonotebook import tqdm
import matplotlib.pyplot as plt
# PyTorch相关
import torch
import torch.nn as nn
import torch.nn.functional as F
# timm库 (PyTorch Image Models)
import timm
# Transformers库
from transformers import (
DistilBertModel,
DistilBertConfig,
DistilBertTokenizer
)
# 图像增强
import albumentations as A
```
## Kaggle数据下载
```python
# 设置Kaggle凭证
os.environ['KAGGLE_USERNAME'] = "your_username"
os.environ['KAGGLE_KEY'] = "your_api_key"
# 下载Flickr 8k数据集
!kaggle datasets download -d adityajn105/flickr8k
!unzip flickr8k.zip
dataset = "8k"
# 或下载Flickr 30k数据集
# !kaggle datasets download -d hsankesara/flickr-image-dataset
# !unzip flickr-image-dataset.zip
# dataset = "30k"
```
## 数据预处理
```python
if dataset == "8k":
df = pd.read_csv("captions.txt")
df['id'] = [id_ for id_ in range(df.shape[0] // 5) for _ in range(5)]
df.to_csv("captions.csv", index=False)
df = pd.read_csv("captions.csv")
image_path = "/content/Images"
captions_path = "/content"
elif dataset == "30k":
df = pd.read_csv("/content/flickr30k_images/results.csv", delimiter="|")
df.columns = ['image', 'caption_number', 'caption']
df['caption'] = df['caption'].str.lstrip()
df['caption_number'] = df['caption_number'].str.lstrip()
df.loc[19999, 'caption_number'] = "4"
df.loc[19999, 'caption'] = "A dog runs across the grass."
ids = [id_ for id_ in range(len(df) // 5) for _ in range(5)]
df['id'] = ids
df.to_csv("captions.csv", index=False)
image_path = "/content/flickr30k_images/flickr30k_images"
captions_path = "/content"
# 查看数据
df.head()
```
## 输出示例
```
image caption id
0 1000268201_693b08cb0e.jpg A child in a pink dress is climbing up a set o... 0
1 1000268201_693b08cb0e.jpg A girl going into a wooden building . 0
2 1000268201_693b08cb0e.jpg A little girl climbing into a wooden playhouse . 0
3 1000268201_693b08cb0e.jpg A little girl climbing the stairs to her playh... 0
4 1000268201_693b08cb0e.jpg A little girl in a pink dress going into a woo... 0
```
## 工具函数
```python
class AvgMeter:
"""用于跟踪和计算平均指标"""
def __init__(self, name="Metric"):
self.name = name
self.reset()
def reset(self):
self.avg, self.sum, self.count = [0] * 3
def update(self, val, count=1):
self.count += count
self.sum += val * count
self.avg = self.sum / self.count
def __repr__(self):
text = f"{self.name}: {self.avg:.4f}"
return text
def get_lr(optimizer):
"""获取当前学习率"""
for param_group in optimizer.param_groups:
return param_group["lr"]
```
|
|
# 配置参数
## CFG配置类
```python
class CFG:
# 调试模式
debug = False
# 路径配置
image_path = image_path
captions_path = captions_path
# 训练参数
batch_size = 32
num_workers = 2
epochs = 4
# 学习率配置
head_lr = 1e-3 # 投影头学习率
image_encoder_lr = 1e-4 # 图像编码器学习率
text_encoder_lr = 1e-5 # 文本编码器学习率
weight_decay = 1e-3
# 学习率调度器
patience = 1
factor = 0.8
# 设备配置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 图像编码器配置
model_name = 'resnet50'
image_embedding = 2048
# 文本编码器配置
text_encoder_model = "distilbert-base-uncased"
text_embedding = 768
text_tokenizer = "distilbert-base-uncased"
max_length = 200
# 预训练配置
pretrained = True # 对图像和文本编码器都启用
trainable = True # 对图像和文本编码器都启用
# 温度参数
temperature = 1.0
# 图像尺寸
size = 224
# 投影头配置
num_projection_layers = 1
projection_dim = 256 # 共享嵌入空间维度
dropout = 0.1
```
## 配置说明
### 学习率策略
| 组件 | 学习率 | 说明 |
|------|-------|------|
| 投影头 | 1e-3 | 从头学习，用较大学习率 |
| 图像编码器 | 1e-4 | 微调预训练模型 |
| 文本编码器 | 1e-5 | 微调预训练模型 |
### 为什么分层学习率？
- **投影头**: 随机初始化，需要快速学习
- **编码器**: 已有预训练权重，需要微调而非破坏
### 嵌入维度
```
图像特征: 2048维 (ResNet50输出)
↓ 投影
共享嵌入: 256维
↑ 投影
文本特征: 768维 (DistilBERT输出)
```
### 温度参数
- 温度=1.0: 无缩放
- 温度>1.0: 使分布更平缓
- 温度<1.0: 使分布更陡峭
## 设备检查
```python
print(f"Using device: {CFG.device}")
print(f"Image size: {CFG.size}x{CFG.size}")
print(f"Batch size: {CFG.batch_size}")
print(f"Projection dim: {CFG.projection_dim}")
```
|
|
# 数据集实现
## CLIPDataset类
```python
class CLIPDataset(torch.utils.data.Dataset):
def __init__(self, image_filenames, captions, tokenizer, transforms):
"""
参数:
image_filenames: 图像文件名列表
captions: 对应的标题列表
tokenizer: HuggingFace分词器
transforms: 图像变换
"""
self.image_filenames = image_filenames
self.captions = list(captions)
# 预先编码所有标题
self.encoded_captions = tokenizer(
list(captions),
padding=True,
truncation=True,
max_length=CFG.max_length
)
self.transforms = transforms
def __getitem__(self, idx):
# 获取预编码的标题
item = {
key: torch.tensor(values[idx])
for key, values in self.encoded_captions.items()
}
# 加载和处理图像
image = cv2.imread(f"{CFG.image_path}/{self.image_filenames[idx]}")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = self.transforms(image=image)['image']
item['image'] = torch.tensor(image).permute(2, 0, 1).float()
item['caption'] = self.captions[idx] # 仅用于可视化
return item
def __len__(self):
return len(self.captions)
```
## 图像变换
```python
def get_transforms(mode="train"):
"""获取图像变换"""
if mode == "train":
return A.Compose([
A.Resize(CFG.size, CFG.size, always_apply=True),
A.Normalize(max_pixel_value=255.0, always_apply=True),
])
else:
return A.Compose([
A.Resize(CFG.size, CFG.size, always_apply=True),
A.Normalize(max_pixel_value=255.0, always_apply=True),
])
```
## 数据加载器
```python
def build_loaders(dataframe, tokenizer, mode):
"""构建数据加载器"""
transforms = get_transforms(mode=mode)
dataset = CLIPDataset(
dataframe["image"].values,
dataframe["caption"].values,
tokenizer=tokenizer,
transforms=transforms,
)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=CFG.batch_size,
num_workers=CFG.num_workers,
shuffle=True if mode == "train" else False,
)
return dataloader
```
## 划分训练集和验证集
```python
def make_train_valid_dfs():
"""划分训练集和验证集"""
dataframe = pd.read_csv(f"{CFG.captions_path}/captions.csv")
max_id = dataframe["id"].max() + 1 if not CFG.debug else 100
# 获取所有图像ID
image_ids = np.arange(0, max_id)
# 随机选择20%作为验证集
np.random.seed(42)
valid_ids = np.random.choice(
image_ids,
size=int(0.2 * len(image_ids)),
replace=False
)
# 划分数据集
train_ids = [id_ for id_ in image_ids if id_ not in valid_ids]
train_dataframe = dataframe[dataframe["id"].isin(train_ids)].reset_index(drop=True)
valid_dataframe = dataframe[dataframe["id"].isin(valid_ids)].reset_index(drop=True)
return train_dataframe, valid_dataframe
```
## 数据集统计
```python
# 使用示例
train_df, valid_df = make_train_valid_dfs()
print(f"训练集大小: {len(train_df)}")
print(f"验证集大小: {len(valid_df)}")
print(f"图像数量: {train_df['id'].nunique()}")
```
## 数据流程图
```
原始CSV文件
↓
读取为DataFrame
↓
划分train/valid
↓
创建CLIPDataset
↓
├─ 预编码文本标题 (Tokenizer)
└─ 加载图像路径
↓
DataLoader
↓
├─ 批量文本 (input_ids, attention_mask)
└─ 批量图像 [batch, 3, 224, 224]
```
|
|
# 模型架构
## 图像编码器
```python
class ImageEncoder(nn.Module):
"""
将图像编码为固定大小的向量
"""
def __init__(
self,
model_name=CFG.model_name,
pretrained=CFG.pretrained,
trainable=CFG.trainable
):
super().__init__()
# 使用timm创建ResNet50
self.model = timm.create_model(
model_name,
pretrained,
num_classes=0, # 移除分类头
global_pool="avg" # 全局平均池化
)
# 设置是否可训练
for p in self.model.parameters():
p.requires_grad = trainable
def forward(self, x):
return self.model(x) # [batch, 2048]
```
## 文本编码器
```python
class TextEncoder(nn.Module):
def __init__(
self,
model_name=CFG.text_encoder_model,
pretrained=CFG.pretrained,
trainable=CFG.trainable
):
super().__init__()
if pretrained:
self.model = DistilBertModel.from_pretrained(model_name)
else:
self.model = DistilBertModel(config=DistilBertConfig())
# 设置是否可训练
for p in self.model.parameters():
p.requires_grad = trainable
# 使用CLS token作为句子表示
self.target_token_idx = 0
def forward(self, input_ids, attention_mask):
output = self.model(
input_ids=input_ids,
attention_mask=attention_mask
)
last_hidden_state = output.last_hidden_state
# 返回CLS token的隐藏状态
return last_hidden_state[:, self.target_token_idx, :]
# [batch, 768]
```
## 投影头
```python
class ProjectionHead(nn.Module):
def __init__(
self,
embedding_dim,
projection_dim=CFG.projection_dim,
dropout=CFG.dropout
):
super().__init__()
self.projection = nn.Linear(embedding_dim, projection_dim)
self.gelu = nn.GELU()
self.fc = nn.Linear(projection_dim, projection_dim)
self.dropout = nn.Dropout(dropout)
self.layer_norm = nn.LayerNorm(projection_dim)
def forward(self, x):
projected = self.projection(x)
x = self.gelu(projected)
x = self.fc(x)
x = self.dropout(x)
x = x + projected # 残差连接
x = self.layer_norm(x)
return x # [batch, 256]
```
## 完整CLIP模型
```python
class CLIPModel(nn.Module):
def __init__(
self,
temperature=CFG.temperature,
image_embedding=CFG.image_embedding,
text_embedding=CFG.text_embedding,
):
super().__init__()
self.image_encoder = ImageEncoder()
self.text_encoder = TextEncoder()
self.image_projection = ProjectionHead(embedding_dim=image_embedding)
self.text_projection = ProjectionHead(embedding_dim=text_embedding)
self.temperature = temperature
def forward(self, batch):
# 获取图像和文本特征
image_features = self.image_encoder(batch["image"])
text_features = self.text_encoder(
input_ids=batch["input_ids"],
attention_mask=batch["attention_mask"]
)
# 获取图像和文本嵌入（相同维度）
image_embeddings = self.image_projection(image_features)
text_embeddings = self.text_projection(text_features)
# 计算损失
logits = (text_embeddings @ image_embeddings.T) / self.temperature
# 计算相似度矩阵作为目标
images_similarity = image_embeddings @ image_embeddings.T
texts_similarity = text_embeddings @ text_embeddings.T
targets = F.softmax(
(images_similarity + texts_similarity) / 2 * self.temperature,
dim=-1
)
# 计算交叉熵损失
texts_loss = cross_entropy(logits, targets, reduction='none')
images_loss = cross_entropy(logits.T, targets.T, reduction='none')
loss = (images_loss + texts_loss) / 2.0
return loss.mean()
def cross_entropy(preds, targets, reduction='none'):
"""自定义交叉熵损失"""
log_softmax = nn.LogSoftmax(dim=-1)
loss = (-targets * log_softmax(preds)).sum(1)
if reduction == "none":
return loss
elif reduction == "mean":
return loss.mean()
```
## 模型架构可视化
```
┌────────────────────────────────────────────────┐
│ CLIP模型完整架构 │
└────────────────────────────────────────────────┘
输入图像 输入文本
[batch, 3, 224, 224] [batch, max_length]
↓ ↓
┌──────────────┐ ┌──────────────┐
│Image Encoder │ │Text Encoder │
│ (ResNet50) │ │(DistilBERT) │
│ timm │ │ Transformers│
└──────┬───────┘ └──────┬───────┘
│ │
↓ ↓
图像特征 文本特征
[batch, 2048] [batch, 768]
│ │
↓ ↓
┌──────────────┐ ┌──────────────┐
│ Projection │ │ Projection │
│ Head │ │ Head │
│ (2048→256) │ │ (768→256) │
└──────┬───────┘ └──────┬───────┘
│ │
↓ ↓
图像嵌入 文本嵌入
[batch, 256] [batch, 256]
│ │
└──────────┬───────────────┘
↓
计算相似度矩阵 / 对比损失
```
|
|
# 训练流程
## 训练一个Epoch
```python
def train_epoch(model, train_loader, optimizer, lr_scheduler, step):
"""训练一个epoch"""
loss_meter = AvgMeter()
tqdm_object = tqdm(train_loader, total=len(train_loader))
for batch in tqdm_object:
# 将数据移到设备
batch = {k: v.to(CFG.device)
for k, v in batch.items() if k != "caption"}
# 前向传播
loss = model(batch)
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 按batch更新学习率
if step == "batch":
lr_scheduler.step()
# 更新损失
count = batch["image"].size(0)
loss_meter.update(loss.item(), count)
tqdm_object.set_postfix(
train_loss=loss_meter.avg,
lr=get_lr(optimizer)
)
return loss_meter
```
## 验证一个Epoch
```python
def valid_epoch(model, valid_loader):
"""验证一个epoch"""
loss_meter = AvgMeter()
tqdm_object = tqdm(valid_loader, total=len(valid_loader))
for batch in tqdm_object:
batch = {k: v.to(CFG.device)
for k, v in batch.items() if k != "caption"}
loss = model(batch)
count = batch["image"].size(0)
loss_meter.update(loss.item(), count)
tqdm_object.set_postfix(valid_loss=loss_meter.avg)
return loss_meter
```
## 主训练函数
```python
def main():
# 准备数据
train_df, valid_df = make_train_valid_dfs()
tokenizer = DistilBertTokenizer.from_pretrained(CFG.text_tokenizer)
train_loader = build_loaders(train_df, tokenizer, mode="train")
valid_loader = build_loaders(valid_df, tokenizer, mode="valid")
# 创建模型
model = CLIPModel().to(CFG.device)
# 配置不同层的学习率
params = [
{
"params": model.image_encoder.parameters(),
"lr": CFG.image_encoder_lr
},
{
"params": model.text_encoder.parameters(),
"lr": CFG.text_encoder_lr
},
{
"params": itertools.chain(
model.image_projection.parameters(),
model.text_projection.parameters()
),
"lr": CFG.head_lr,
"weight_decay": CFG.weight_decay
}
]
# 优化器
optimizer = torch.optim.AdamW(params, weight_decay=0.)
# 学习率调度器
lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer,
mode="min",
patience=CFG.patience,
factor=CFG.factor
)
step = "epoch"
# 训练循环
best_loss = float('inf')
for epoch in range(CFG.epochs):
print(f"Epoch: {epoch + 1}")
# 训练
model.train()
train_loss = train_epoch(
model, train_loader, optimizer, lr_scheduler, step
)
# 验证
model.eval()
with torch.no_grad():
valid_loss = valid_epoch(model, valid_loader)
# 保存最佳模型
if valid_loss.avg < best_loss:
best_loss = valid_loss.avg
torch.save(model.state_dict(), "best.pt")
print("Saved Best Model!")
# 更新学习率
lr_scheduler.step(valid_loss.avg)
```
## 运行训练
```python
# 开始训练
main()
```
## 训练输出示例
```
Epoch: 1
train_loss: 3.2456, lr: 0.0001: 100%|████████| 1012/1012 [08:23<00:00, 2.01it/s]
valid_loss: 2.8765: 100%|████████| 253/253 [01:15<00:00, 3.36it/s]
Saved Best Model!
Epoch: 2
train_loss: 2.5432, lr: 0.0001: 100%|████████| 1012/1012 [08:21<00:00, 2.02it/s]
valid_loss: 2.6543: 100%|████████| 253/253 [01:15<00:00, 3.36it/s]
Saved Best Model!
...
```
## 训练技巧
| 技巧 | 说明 |
|------|------|
| **分层学习率** | 编码器小学习率，投影头大学习率 |
| **权重衰减** | 防止过拟合 (1e-3) |
| **学习率调度** | ReduceLROnPlateau自适应调整 |
| **梯度裁剪** | 防止梯度爆炸 |
| **混合精度** | 使用AMP加速训练 |
|
|
# 推理应用
## 获取图像嵌入
```python
def get_image_embeddings(valid_df, model_path):
"""
获取验证集所有图像的嵌入
返回:
model: 加载的CLIP模型
image_embeddings: 所有图像的嵌入 [valid_set_size, 256]
"""
tokenizer = DistilBertTokenizer.from_pretrained(CFG.text_tokenizer)
valid_loader = build_loaders(valid_df, tokenizer, mode="valid")
model = CLIPModel().to(CFG.device)
model.load_state_dict(torch.load(model_path, map_location=CFG.device))
model.eval()
valid_image_embeddings = []
with torch.no_grad():
for batch in tqdm(valid_loader):
image_features = model.image_encoder(batch["image"].to(CFG.device))
image_embeddings = model.image_projection(image_features)
valid_image_embeddings.append(image_embeddings)
return model, torch.cat(valid_image_embeddings)
```
## 查找匹配图像
```python
import matplotlib.pyplot as plt
def find_matches(model, image_embeddings, query, image_filenames, n=9):
"""
根据文本查询查找匹配的图像
参数:
model: CLIP模型
image_embeddings: 预计算的图像嵌入
query: 查询文本
image_filenames: 图像文件名列表
n: 显示的图像数量
"""
tokenizer = DistilBertTokenizer.from_pretrained(CFG.text_tokenizer)
# 编码查询文本
encoded_query = tokenizer([query])
batch = {
key: torch.tensor(values).to(CFG.device)
for key, values in encoded_query.items()
}
with torch.no_grad():
# 获取文本特征
text_features = model.text_encoder(
input_ids=batch["input_ids"],
attention_mask=batch["attention_mask"]
)
text_embeddings = model.text_projection(text_features)
# 归一化
image_embeddings_n = F.normalize(image_embeddings, p=2, dim=-1)
text_embeddings_n = F.normalize(text_embeddings, p=2, dim=-1)
# 计算点积相似度
dot_similarity = text_embeddings_n @ image_embeddings_n.T
# 获取top-k匹配
values, indices = torch.topk(dot_similarity.squeeze(0), n * 5)
matches = [image_filenames[idx] for idx in indices[::5]]
# 显示结果
_, axes = plt.subplots(3, 3, figsize=(10, 10))
for match, ax in zip(matches, axes.flatten()):
image = cv2.imread(f"{CFG.image_path}/{match}")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
ax.imshow(image)
ax.axis("off")
plt.show()
```
## 使用示例
```python
# 获取验证集图像嵌入
_, valid_df = make_train_valid_dfs()
model, image_embeddings = get_image_embeddings(valid_df, "best.pt")
# 查询1: 狗在草地上
find_matches(
model,
image_embeddings,
query="dogs on the grass",
image_filenames=valid_df['image'].values,
n=9
)
# 查询2: 男孩玩滑板
find_matches(
model,
image_embeddings,
query="a boy jumping with skateboard",
image_filenames=valid_df['image'].values,
n=9
)
```
## 零样本分类
```python
def zero_shot_classifier(model, image, class_prompts):
"""
零样本图像分类
参数:
model: CLIP模型
image: 输入图像
class_prompts: 类别提示列表
返回:
predicted_class: 预测的类别
"""
tokenizer = DistilBertTokenizer.from_pretrained(CFG.text_tokenizer)
# 编码类别提示
encoded_prompts = tokenizer(class_prompts)
batch = {
key: torch.tensor(values).to(CFG.device)
for key, values in encoded_prompts.items()
}
with torch.no_grad():
# 获取文本嵌入
text_features = model.text_encoder(
input_ids=batch["input_ids"],
attention_mask=batch["attention_mask"]
)
text_embeddings = model.text_projection(text_features)
# 获取图像嵌入
image_features = model.image_encoder(image)
image_embedding = model.image_projection(image_features)
# 归一化
text_embeddings_n = F.normalize(text_embeddings, p=2, dim=-1)
image_embedding_n = F.normalize(image_embedding, p=2, dim=-1)
# 计算相似度
similarity = image_embedding_n @ text_embeddings_n.T
# 返回最匹配的类别
predicted_idx = similarity.argmax()
return class_prompts[predicted_idx]
# 使用示例
class_prompts = [
"a photo of a dog",
"a photo of a cat",
"a photo of a bird",
"a photo of a car"
]
# 预测
# predicted = zero_shot_classifier(model, test_image, class_prompts)
```
## 应用场景总结
| 应用 | 说明 |
|------|------|
| **图像检索** | 根据文本描述检索图像 |
| **零样本分类** | 无需训练即可分类新类别 |
| **自动标注** | 为图像生成描述 |
| **内容审核** | 检测不当内容 |
| **推荐系统** | 基于视觉相似度推荐 |
|
|
# CLIP源码目录结构
https://github.com/openai/CLIP/tree/main
OpenAI CLIP代码库采用模块化设计，主要文件如下：
```
CLIP-main/
├── clip/
│ ├── __init__.py # 包初始化，导出主要API
│ ├── clip.py # 用户接口：模型加载、图像预处理、文本分词
│ ├── model.py # 核心模型定义：ResNet/ViT、Transformer
│ └── simple_tokenizer.py # BPE文本分词器
├── setup.py # 包安装配置
├── hubconf.py # PyTorch Hub配置
└── tests/ # 测试文件
```
### 代码层次结构
```
用户层 (clip.py)
├── load() - 加载预训练模型
├── tokenize() - 文本分词
└── _transform() - 图像预处理变换
↓
模型层 (model.py)
├── CLIP - 主模型类
├── VisionTransformer / ModifiedResNet - 图像编码器
├── Transformer - 文本编码器
├── AttentionPool2d - 注意力池化
└── ResidualAttentionBlock - Transformer块
↓
工具层 (simple_tokenizer.py)
├── SimpleTokenizer - BPE分词器
├── bpe() - 字节对编码
└── encode/decode - 编解码
```
|
|
## clip.py 详解
### 导入依赖 (第1-28行)
```python
# 标准库
import hashlib # 用于SHA256校验
import os # 文件路径操作
import urllib # 模型下载
# 版本管理
from packaging import version
from typing import Union, List
# 深度学习框架
import torch
from PIL import Image
from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
from tqdm import tqdm
# CLIP内部模块
from .model import build_model
from .simple_tokenizer import SimpleTokenizer as _Tokenizer
```
**关键点说明：**
- `hashlib`：下载模型后验证完整性，确保模型未被篡改
- `torchvision.transforms`：图像预处理管道
- `build_model`：从state_dict构建模型实例
### 模型配置字典 (第30-40行)
```python
_MODELS = {
"RN50": "https://openaipublic.azureedge.net/clip/models/...",
"RN101": "...",
"RN50x4": "...",
"RN50x16": "...",
"RN50x64": "...",
"ViT-B/32": "...",
"ViT-B/16": "...",
"ViT-L/14": "...",
"ViT-L/14@336px": "...",
}
```
**模型命名规则：**
- `RNxx`：ResNet架构
- `RN50`：标准ResNet-50
- `RN50x4`：宽度放大4倍
- `RN50x16`：宽度放大16倍
- `RN50x64`：宽度放大64倍
- `ViT-B/32`：Vision Transformer，Base大小，32x32像素patch
- `ViT-L/14`：Vision Transformer，Large大小，14x14像素patch
### 模型下载函数 (第43-72行)
```python
def _download(url: str, root: str):
# 1. 创建目录
os.makedirs(root, exist_ok=True)
# 2. 从URL提取期望的SHA256
expected_sha256 = url.split("/")[-2]
# 3. 检查文件是否已存在且校验通过
if os.path.isfile(download_target):
if hashlib.sha256(...).hexdigest() == expected_sha256:
return download_target
# 4. 下载并显示进度条
with urllib.request.urlopen(url) as source:
with tqdm(...) as loop:
while True:
buffer = source.read(8192) # 8KB块读取
if not buffer:
break
output.write(buffer)
loop.update(len(buffer))
# 5. 验证下载的文件
if hashlib.sha256(...).hexdigest() != expected_sha256:
raise RuntimeError("校验和不匹配")
```
**下载流程图：**
```
开始 → 检查本地文件 → 校验SHA256
↓ ↓
已存在且有效 需下载
↓ ↓
直接返回 下载 → 校验 → 返回
↓
校验失败 → 报错
```
### 图像预处理 (第79-86行)
```python
def _transform(n_px):
return Compose([
Resize(n_px, interpolation=BICUBIC), # 双三次插值缩放
CenterCrop(n_px), # 中心裁剪
_convert_image_to_rgb, # 转RGB
ToTensor(), # 转张量 [0,1]
Normalize((0.48145466, 0.4578275, 0.40821073), # ImageNet均值
(0.26862954, 0.26130258, 0.27577711)) # ImageNet标准差
])
```
**预处理步骤详解：**
| 步骤 | 输入 | 输出 | 说明 |
|------|------|------|------|
| Resize | 原图 | 短边=n_px | 保持宽高比 |
| CenterCrop | 矩形 | n_px×n_px | 中心裁剪为正方形 |
| ToTensor | PIL Image | Tensor[0,1] | 归一化到[0,1] |
| Normalize | Tensor[0,1] | Tensor[-1,1] | 减均值除标准差 |
### 文本分词 (第205-245行)
```python
def tokenize(texts: Union[str, List[str]],
context_length: int = 77,
truncate: bool = False):
"""
将文本转换为CLIP输入token
输出格式：
<|startoftext|> token1 token2 ... <|endoftext|>
|
|
## model.py 详解
### 文件结构概览
```python
# 1. ResNet相关组件
class Bottleneck(nn.Module) # ResNet瓶颈块
class AttentionPool2d(nn.Module) # 注意力池化层
class ModifiedResNet(nn.Module) # 修改版ResNet编码器
# 2. Transformer相关组件
class LayerNorm(nn.LayerNorm) # LayerNorm（支持fp16）
class QuickGELU(nn.Module) # 快速GELU激活
class ResidualAttentionBlock(nn.Module) # Transformer残差块
class Transformer(nn.Module) # Transformer编码器
# 3. Vision Transformer
class VisionTransformer(nn.Module) # ViT图像编码器
# 4. CLIP主模型
class CLIP(nn.Module) # 完整CLIP模型
# 5. 工具函数
def convert_weights(model) # 转换为fp16
def build_model(state_dict) # 从权重构建模型
```
### Bottleneck类 (第10-55行)
```python
class Bottleneck(nn.Module):
"""ResNet瓶颈块，扩展因子=4"""
expansion = 4
def __init__(self, inplanes, planes, stride=1):
super().__init__()
# 1×1卷积降维
self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
self.bn1 = nn.BatchNorm2d(planes)
self.relu1 = nn.ReLU(inplace=True)
# 3×3卷积（主操作）
self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes)
self.relu2 = nn.ReLU(inplace=True)
# 步幅>1时使用avgpool下采样
self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
# 1×1卷积升维（4倍扩展）
self.conv3 = nn.Conv2d(planes, planes * 4, 1, bias=False)
self.bn3 = nn.BatchNorm2d(planes * 4)
self.relu3 = nn.ReLU(inplace=True)
# 残差连接的降采样
self.downsample = None
if stride > 1 or inplanes != planes * 4:
self.downsample = nn.Sequential([
nn.AvgPool2d(stride),
nn.Conv2d(inplanes, planes * 4, 1, bias=False),
nn.BatchNorm2d(planes * 4)
])
def forward(self, x):
identity = self.downsample(x) if self.downsample else x
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))
out += identity # 残差连接
return self.relu3(out)
```
**Bottleneck结构图：**
```
输入 → 1×1conv → BN → ReLU → 3×3conv → BN → ReLU → [AvgPool] → 1×1conv → BN → + → ReLU
↑ ↑
（步幅>1时） 残差连接
```
### AttentionPool2d类 (第58-91行)
```python
class AttentionPool2d(nn.Module):
"""用多头注意力代替传统平均池化"""
def __init__(self, spacial_dim, embed_dim, num_heads, output_dim=None):
super().__init__()
# 位置编码：HW+1个位置（+1是CLS token）
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)
# 添加CLS token（空间平均）
x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)
# 添加位置编码
x = x + self.positional_embedding[:, None, :]
# 多头注意力：CLS作为query，所有位置作为key/value
x, _ = F.multi_head_attention_forward(
query=x[:1], # CLS token查询
key=x, # 所有位置作为key
value=x, # 所有位置作为value
...
)
return x.squeeze(0)
```
**AttentionPool原理：**
```
特征图 → 展平 → 添加CLS → 加位置编码 → 多头注意力 → 输出
↑
全局平均池化作为CLS token
```
### ModifiedResNet类 (第94-154行)
```python
class ModifiedResNet(nn.Module):
"""
与torchvision ResNet的区别：
1. 3层stem（而非1层）
2. 使用avgpool而非maxpool
3. 步幅>1的卷积前加avgpool（抗混叠）
4. 最终用注意力池化替代平均池化
"""
def __init__(self, layers, output_dim, heads, input_resolution=224, width=64):
super().__init__()
# 3层stem：逐步下采样
self.conv1 = nn.Conv2d(3, width//2, 3, stride=2, padding=1) # 112
self.conv2 = nn.Conv2d(width//2, width//2, 3, padding=1) # 112
self.conv3 = nn.Conv2d(width//2, width, 3, padding=1) # 112
self.avgpool = nn.AvgPool2d(2) # 56
# 4个残差层
self.layer1 = self._make_layer(width, layers[0]) # 56
self.layer2 = self._make_layer(width*2, layers[1], 2) # 28
self.layer3 = self._make_layer(width*4, layers[2], 2) # 14
self.layer4 = self._make_layer(width*8, layers[3], 2) # 7
# 注意力池化：7×7→输出维度
self.attnpool = AttentionPool2d(7, width*32, heads, output_dim)
```
**ModifiedResNet结构：**
```
输入(224×224)
↓
stem(3层conv + avgpool) → 56×56×width
↓
layer1 (×1) → 56×56×4×width
↓
layer2 (×2, stride=2) → 28×28×8×width
↓
layer3 (×4, stride=2) → 14×14×16×width
↓
layer4 (×8, stride=2) → 7×7×32×width
↓
AttentionPool2d → output_dim
```
### VisionTransformer类 (第206-240行)
```python
class VisionTransformer(nn.Module):
"""Vision Transformer图像编码器"""
def __init__(self, input_resolution, patch_size, width, layers, heads, output_dim):
super().__init__()
# Patch嵌入：conv2d将patch展平为向量
self.conv1 = nn.Conv2d(3, width, kernel_size=patch_size, stride=patch_size)
# CLS token和位置编码
scale = width ** -0.5
self.class_embedding = nn.Parameter(scale * torch.randn(width))
grid_size = input_resolution // patch_size
self.positional_embedding = nn.Parameter(
scale * torch.randn(grid_size ** 2 + 1, width)
)
# 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):
# 1. 分patch：N×3×H×W → N×width×grid×grid
x = self.conv1(x)
# 2. 展平：N×width×grid×grid → N×width×grid²
x = x.reshape(x.shape[0], x.shape[1], -1)
# 3. 转置：N×width×grid² → N×grid²×width
x = x.permute(0, 2, 1)
# 4. 添加CLS token：N×grid²×width → N×(grid²+1)×width
x = torch.cat([
self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], ...),
x
], dim=1)
# 5. 加位置编码
x = x + self.positional_embedding.to(x.dtype)
# 6. LayerNorm
x = self.ln_pre(x)
# 7. Transformer (需要LND格式)
x = x.permute(1, 0, 2) # NLD → LND
x = self.transformer(x)
x = x.permute(1, 0, 2) # LND → NLD
# 8. 取CLS token输出
x = self.ln_post(x[:, 0, :])
# 9. 投影到输出维度
return x @ self.proj
```
**ViT处理流程：**
```
图像 → Patch嵌入 → 加CLS/位置编码 → Transformer → 取CLS → 投影
[N×3×224×224] [N×(7²+1)×768] [N×768] [N×512]
```
### Transformer相关类
**QuickGELU (第166-168行)：**
```python
class QuickGELU(nn.Module):
"""快速GELU近似：x * sigmoid(1.702*x)"""
def forward(self, x):
return x * torch.sigmoid(1.702 * x)
```
**ResidualAttentionBlock (第171-192行)：**
```python
class ResidualAttentionBlock(nn.Module):
"""Transformer残差块：Pre-LN结构"""
def __init__(self, d_model, n_head, attn_mask=None):
super().__init__()
self.attn = nn.MultiheadAttention(d_model, n_head)
self.ln_1 = LayerNorm(d_model) # Pre-LN
# MLP
self.mlp = nn.Sequential([
nn.Linear(d_model, d_model * 4),
QuickGELU(),
nn.Linear(d_model * 4, d_model)
])
self.ln_2 = LayerNorm(d_model) # Pre-LN
self.attn_mask = attn_mask
def forward(self, x):
# Pre-LN + 残差
x = x + self.attention(self.ln_1(x))
x = x + self.mlp(self.ln_2(x))
return x
```
**Transformer结构图：**
```
输入 → LN → MultiHeadAttention → + → LN → MLP → + → 输出
↑ ↑ ↑ ↑
Pre-LN 残差 Pre-LN 残差
```
### CLIP主类 (第243-372行)
```python
class CLIP(nn.Module):
def __init__(self, embed_dim, image_resolution, vision_layers,
vision_width, vision_patch_size, context_length,
vocab_size, transformer_width, transformer_heads,
transformer_layers):
super().__init__()
# 图像编码器：ResNet或ViT
if isinstance(vision_layers, tuple):
self.visual = ModifiedResNet(...)
else:
self.visual = VisionTransformer(...)
# 文本编码器：Transformer
self.transformer = Transformer(
width=transformer_width,
layers=transformer_layers,
heads=transformer_heads,
attn_mask=self.build_attention_mask() # 因果mask
)
# 文本嵌入
self.token_embedding = nn.Embedding(vocab_size, transformer_width)
self.positional_embedding = nn.Parameter(
torch.empty(context_length, transformer_width)
)
self.ln_final = LayerNorm(transformer_width)
# 投影到共享空间
self.text_projection = nn.Parameter(
torch.empty(transformer_width, embed_dim)
)
# 温度参数（初始化为1/0.07的log）
self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
def encode_image(self, image):
"""图像编码"""
return self.visual(image.type(self.dtype))
def encode_text(self, text):
"""文本编码"""
# 1. token嵌入
x = self.token_embedding(text).type(self.dtype)
# 2. 加位置编码
x = x + self.positional_embedding.type(self.dtype)
# 3. Transformer
x = x.permute(1, 0, 2) # NLD → LND
x = self.transformer(x)
x = x.permute(1, 0, 2) # LND → NLD
# 4. LayerNorm
x = self.ln_final(x).type(self.dtype)
# 5. 取EOT token的特征
x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)]
# 6. 投影到共享空间
return x @ self.text_projection
def forward(self, image, text):
"""计算图像-文本相似度"""
# 1. 编码
image_features = self.encode_image(image)
text_features = self.encode_text(text)
# 2. L2归一化
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()
logits_per_image = logit_scale * (image_features @ text_features.T)
logits_per_text = logits_per_image.T
return logits_per_image, logits_per_text
```
**CLIP forward流程：**
```
图像 + 文本
↓
图像编码器 → 图像特征(768) → 归一化
文本编码器 → 文本特征(768) → 归一化
↓
计算余弦相似度矩阵 (×温度参数)
↓
返回 logits_per_image (N×N)
logits_per_text (N×N)
```
### build_model函数 (第399-436行)
```python
def build_model(state_dict):
"""从state_dict构建模型"""
# 1. 检测是ViT还是ResNet
vit = "visual.proj" in state_dict
if vit:
# 解析ViT参数
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:
# 解析ResNet参数
counts = [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]
image_resolution = output_width * 32
# 2. 解析文本编码器参数
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")))
# 3. 构建模型
model = CLIP(embed_dim, image_resolution, vision_layers, vision_width,
vision_patch_size, context_length, vocab_size,
transformer_width, transformer_heads, transformer_layers)
# 4. 转换权重并加载
convert_weights(model)
model.load_state_dict(state_dict)
return model.eval()
```
|
|
## simple_tokenizer.py 详解
### bytes_to_unicode函数 (第16-35行)
```python
@lru_cache()
def bytes_to_unicode():
"""
字节到Unicode的映射表
目的：BPE在Unicode上工作，需要避免大量UNK token
将256个字节映射到可Unicode字符范围
"""
# 1. 可打印ASCII字符
bs = list(range(ord("!"), ord("~")+1)) + \
list(range(ord("¡"), ord("¬")+1)) + \
list(range(ord("®"), ord("ÿ")+1))
# 2. 其他字节映射到高Unicode范围
cs = bs[:]
n = 0
for b in range(2**8):
if b not in bs:
bs.append(b)
cs.append(2**8 + n)
n += 1
cs = [chr(n) for n in cs]
return dict(zip(bs, cs))
```
**字节映射示例：**
```
字节值 → Unicode字符
65-126 → 直接可见字符（A-Z, a-z, 0-9, 标点）
其他 → 高Unicode字符（Ā, ā, Ă...）
```
### SimpleTokenizer类 (第62-132行)
```python
class SimpleTokenizer:
def __init__(self, bpe_path=default_bpe()):
# 1. 字节编解码器
self.byte_encoder = bytes_to_unicode()
self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
# 2. 加载BPE合并规则
merges = gzip.open(bpe_path).read().decode("utf-8").split('\n')
merges = merges[1:49152-256-2+1] # 跳过首行和特殊token
merges = [tuple(merge.split()) for merge in merges]
# 3. 构建词汇表
vocab = list(bytes_to_unicode().values())
vocab = vocab + [v+'' for v in vocab] # 词尾标记
for merge in merges:
vocab.append(''.join(merge))
vocab.extend(['<|startoftext|>', '<|endoftext|>'])
# 4. 编解码映射
self.encoder = dict(zip(vocab, range(len(vocab))))
self.decoder = {v: k for k, v in self.encoder.items()}
self.bpe_ranks = dict(zip(merges, range(len(merges))))
# 5. 正则表达式（分词用）
self.pat = re.compile(
r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|
[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""",
re.IGNORECASE
)
```
**词汇表构成：**
```
基础：256个Unicode字符
扩展：每个字符+词尾标记（256个）
BPE：48888个合并规则
特殊：<|startoftext|>, <|endoftext|>
总计：49408个token
```
### BPE编码 (第80-119行)
```python
def bpe(self, token):
"""字节对编码：贪心合并高频字符对"""
if token in self.cache:
return self.cache[token]
# 1. 添加词尾标记
word = tuple(token[:-1]) + (token[-1] + '',)
# 2. 获取所有字符对
pairs = get_pairs(word)
if not pairs:
return token + ''
# 3. 贪心合并
while True:
# 找到最高优先级的字符对
bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float('inf')))
if bigram not in self.bpe_ranks:
break
first, second = bigram
new_word = []
i = 0
while i < len(word):
try:
j = word.index(first, i)
new_word.extend(word[i:j])
i = j
except:
new_word.extend(word[i:])
break
# 合并匹配的字符对
if word[i] == first and i < len(word)-1 and word[i+1] == second:
new_word.append(first + second)
i += 2
else:
new_word.append(word[i])
i += 1
word = tuple(new_word)
if len(word) == 1:
break
else:
pairs = get_pairs(word)
return ' '.join(word)
```
**BPE示例：**
```
输入: "hug"
1. 初始: ['h', 'u', 'g']
2. 字符对: ('h','u'), ('u','g')
3. 假设('u','g')优先级最高
4. 合并: ['h', 'ug']
输出: "h ug"
```
### encode函数 (第121-127行)
```python
def encode(self, text):
"""文本→token ID序列"""
bpe_tokens = []
# 1. 清理和转小写
text = whitespace_clean(basic_clean(text)).lower()
# 2. 正则分词
for token in re.findall(self.pat, text):
# 3. 字节→Unicode
token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
# 4. BPE编码
for bpe_token in self.bpe(token).split(' '):
bpe_tokens.append(self.encoder[bpe_token])
return bpe_tokens
```
**encode完整示例：**
```
输入: "Hello World!"
1. 清理: "hello world!"
2. 分词: ["hello", "world", "!"]
3. 字节映射: ["Ħėļļō", "ŵōřļđ", "!"]
4. BPE: ["Ħ", "ė", "ļļ", "ō", "ŵ", "ō", "ř", "ļ", "đ", "!"]
5. 编码: [66, 76, 399, 77, 95, 77, 96, 399, 97, 0]
```
### decode函数 (第129-132行)
```python
def decode(self, tokens):
"""token ID序列→文本"""
# 1. ID→字符
text = ''.join([self.decoder[token] for token in tokens])
# 2. Unicode→字节
text = bytearray([self.byte_decoder[c] for c in text])
# 3. 字节→UTF-8字符串
text = text.decode('utf-8', errors="replace").replace('', ' ')
return text
```
|
|
## CLIP核心工作流程
### 完整推理流程
```python
import torch
import clip
from PIL import Image
# 1. 加载模型
device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-B/32", device=device)
# 2. 准备输入
image = preprocess(Image.open("dog.jpg")).unsqueeze(0).to(device)
text = clip.tokenize(["a dog", "a cat", "a car"]).to(device)
# 3. 前向传播
with torch.no_grad():
image_features = model.encode_image(image) # [1, 512]
text_features = model.encode_text(text) # [3, 512]
# 4. 计算相似度
logits_per_image, logits_per_text = model(image, text)
probs = logits_per_image.softmax(dim=-1).cpu().numpy()
print("Label probs:", probs) # [dog: 0.99, cat: 0.01, car: 0.00]
```
### 数据流动图
```
用户输入
↓
┌─────────────┬─────────────┐
│ │ │
图像 (224×224) 文本 ("a dog")
│ │ │
↓ ↓ │
preprocess() tokenize() │
│ │ │
↓ ↓ │
Tensor[3,224,224] Tensor[1,77] │
│ │ │
↓ ↓ │
┌─────────────┴─────────────┐
│ │
│ CLIP模型 │
│ ┌──────────────────────┐ │
│ │ 图像编码器 │ │
│ │ (ViT/ResNet) │ │
│ │ → 512维向量 │ │
│ └──────────────────────┘ │
│ │
│ ┌──────────────────────┐ │
│ │ 文本编码器 │ │
│ │ (Transformer) │ │
│ │ → 512维向量 │ │
│ └──────────────────────┘ │
│ │
│ 余弦相似度 × 温度参数 │
└─────────────────────────────┘
↓
相似度矩阵 [1, 3]
↓
softmax → 概率分布
```
### 训练流程（对比学习）
```python
# 批次数据
images = preprocess_batch([img1, img2, ..., imgN]) # [N, 3, 224, 224]
texts = tokenize(["caption 1", "caption 2", ...]) # [N, 77]
# 前向传播
logits_per_image, logits_per_text = model(images, texts)
# 真实标签：对角线为匹配对
labels = torch.arange(N).to(device)
# 损失：交叉熵
loss_i = F.cross_entropy(logits_per_image, labels) # 图像→文本
loss_t = F.cross_entropy(logits_per_text, labels) # 文本→图像
loss = (loss_i + loss_t) / 2
# 反向传播
loss.backward()
optimizer.step()
```
### 损失计算详解
```
相似度矩阵（batch_size=3）：
文本1 文本2 文本3
图像1 [ 10.5 -5.2 -8.1 ] ← 真实匹配
图像2 [ -4.8 9.7 -6.3 ]
图像3 [ -7.2 -5.9 8.8 ] ← 真实匹配
对于图像1：
- 正样本对：(图像1, 文本1) → logit=10.5
- 负样本对：(图像1, 文本2/3) → logit=-5.2, -8.1
- 损失：-log(exp(10.5) / (exp(10.5)+exp(-5.2)+exp(-8.1)))
```
### 零样本分类
```python
# 图像分类提示模板
class_prompts = [
"a photo of a dog",
"a photo of a cat",
"a photo of a car",
"a photo of a bird"
]
# 编码文本
text_tokens = clip.tokenize(class_prompts).to(device)
# 推理
with torch.no_grad():
image_features = model.encode_image(image)
text_features = model.encode_text(text_tokens)
# 计算相似度
similarity = (image_features @ text_features.T).squeeze(0)
probs = similarity.softmax(dim=-1)
# 获取预测类别
pred_idx = probs.argmax().item()
print(f"Predicted: {class_prompts[pred_idx]}")
```
### 关键参数说明
| 参数 | ViT-B/32 | ViT-B/16 | ViT-L/14 |
|------|----------|----------|----------|
| 图像分辨率 | 224 | 224 | 224/336 |
| Patch大小 | 32×32 | 16×16 | 14×14 |
| Embedding维度 | 768 | 768 | 1024 |
| Transformer层数 | 12 | 12 | 24 |
| Attention头数 | 12 | 12 | 16 |
| 输出维度 | 512 | 512 | 768 |
| 参数量 | 150M | 86M | 304M |
### 温度参数作用
```python
# 温度参数控制softmax的锐度
logit_scale = model.logit_scale.exp() # 默认约1/0.07 ≈ 14.29
# 高温度 → 较低的相似度 → 更均匀的分布
# 低温度 → 较高的相似度 → 更锐利的分布
```
**温度缩放效果：**
```
未缩放相似度: [2.0, 1.0, 0.5]
温度=1: softmax → [0.66, 0.24, 0.10]
温度=10: softmax → [0.45, 0.32, 0.23] ← 更均匀
温度=0.1: softmax → [0.92, 0.07, 0.01] ← 更锐利
```
|
|
# OpenCLIP项目概述
https://github.com/mlfoundations/open_clip
## 项目简介
OpenCLIP是MLFoundations团队开源的CLIP实现，支持大规模对比语言-图像预训练。相比OpenAI原版CLIP，OpenCLIP提供了更丰富的模型变体、训练数据和配置选项。
## 核心特性
### 支持的模型类型
| 模型类型 | 描述 | 配置前缀 |
|---------|------|---------|
| ViT | Vision Transformer | `ViT-B-*`, `ViT-L-*`, `ViT-H-*`, `ViT-g-*` |
| ConvNeXt | 现代CNN架构 | `convnext_*` |
| ResNet | 经典残差网络 | `RN50*` |
| SigLIP | Sigmoid Loss CLIP | `ViT-SO-*-SigLIP` |
| CoCa | Contrastive Captioners | `coca_*` |
| CLAP | Contrastive Language-Audio Pre-training | `clap_*` |
| NaFlex | 灵活分辨率/长宽比 | `naflex*` |
### 支持的模态
```
图像-文本 (CLIP) ─→ 对比学习
音频-文本 (CLAP) ─→ 对比学习
生成式 (GenLIP/GenLAP) ─→ 前缀LM + 打包训练
```
## 目录结构
```
open_clip-main/
├── src/
│ ├── open_clip/ # 核心推理库
│ │ ├── __init__.py # 导出API
│ │ ├── factory.py # 模型工厂（加载/创建）
│ │ ├── model.py # CLIP模型定义
│ │ ├── modified_resnet.py # ResNet图像编码器
│ │ ├── transformer.py # Transformer组件
│ │ ├── timm_model.py # timm图像编码器封装
│ │ ├── hf_model.py # HuggingFace文本编码器
│ │ ├── tokenizer.py # 文本分词器
│ │ ├── transform.py # 图像预处理
│ │ ├── loss.py # 损失函数
│ │ ├── pretrained.py # 预训练权重管理
│ │ ├── task.py # 训练任务封装
│ │ ├── constants.py # 常量定义
│ │ ├── coca_model.py # CoCa模型
│ │ ├── clap_model.py # CLAP模型
│ │ └── model_configs/ # 模型配置JSON
│ │
│ └── open_clip_train/ # 训练相关代码
│ ├── main.py # 训练入口
│ ├── train.py # 训练循环
│ ├── data.py # 数据加载
│ ├── distributed.py # 分布式训练
│ ├── optim.py # 优化器配置
│ ├── scheduler.py # 学习率调度
│ ├── params.py # 参数解析
│ ├── zero_shot.py # 零样本评估
│ ├── metrics.py # 评估指标
│ ├── profiler.py # 性能分析
│ └── legacy_*.py # 兼容旧版代码
│
├── scripts/ # 辅助脚本
├── tests/ # 单元测试
└── docs/ # 文档
```
## 代码架构层次
```
┌─────────────────────────────────────────────────────────────────┐
│ 用户API层 │
│ create_model() / create_model_and_transforms() / get_tokenizer()│
└─────────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────────┐
│ 工厂层 (factory.py) │
│ 模型创建 | 权重加载 | 配置解析 | 分词器解析 │
└─────────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────────┐
│ 模型层 (model.py) │
│ CLIP | CustomTextCLIP | CoCa | CLAP | VisionTransformer │
└─────────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────────┐
│ 组件层 │
│ Transformer | ResNet | Attention | LayerNorm | 损失函数 │
└─────────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────────┐
│ 训练层 (train.py + task.py) │
│ TrainingTask | 前向传播 | 反向传播 | 优化器 | 调度器 │
└─────────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────────┐
│ 数据层 (data.py) │
│ WebDataset | CsvDataset | ImageFolder | 数据增强 │
└─────────────────────────────────────────────────────────────────┘
```
## 与OpenAI CLIP的对比
| 特性 | OpenAI CLIP | OpenCLIP |
|------|-------------|----------|
| 模型类型 | ViT + ResNet | ViT + ResNet + ConvNeXt + SigLIP + CoCa |
| 训练数据 | WIT (400M) | LAION-2B, DataComp-1B等 |
| 预训练权重 | OpenAI服务器 | 多个来源 + 本地路径 + HuggingFace |
| 训练代码 | 简单脚本 | 完整训练框架 |
| 分布式 | 基础DDP | DDP + FSDP2 |
| 精度支持 | FP16/FP32 | FP16/FP32/BF16 + AMP |
| 扩展性 | 有限 | 高度模块化 |
## 依赖关系
```
torch >= 2.6
torchvision
timm (可选，用于图像编码器)
transformers (可选，用于文本编码器)
webdataset (训练时)
ftfy (文本清理)
regex (分词器)
```
## 模型配置系统
OpenCLIP使用JSON配置文件定义模型架构：
```json
{
"embed_dim": 512,
"vision_cfg": {
"image_size": 224,
"layers": 12,
"width": 768,
"patch_size": 32
},
"text_cfg": {
"context_length": 77,
"vocab_size": 49408,
"width": 512,
"layers": 12,
"heads": 8
}
}
```
**支持的配置来源：**
1. 内置配置 (model_configs/*.json)
2. HuggingFace Hub (hf-hub:org/model)
3. 本地目录 (local-dir:/path)
## 预训练权重标签
```python
# 示例权重标签
'laion2b_s34b_b79k' # LAION-2B, 34B样本, 79k批次
'datacomp_xl_s13b_b90k' # DataComp-XL, 13B样本, 90k批次
'openai' # OpenAI原始权重
'redisclosed' # 重新发布的权重
```
|
|
## factory.py 详解
### 模型配置注册系统 (第46-72行)
```python
_MODEL_CONFIG_PATHS = [Path(__file__).parent / "model_configs/"]
_MODEL_CONFIGS = {} # 模型名 → 配置字典
def _rescan_model_configs():
"""扫描model_configs目录，加载所有JSON配置"""
global _MODEL_CONFIGS
config_files = []
for config_path in _MODEL_CONFIG_PATHS:
if config_path.is_file() and config_path.suffix == '.json':
config_files.append(config_path)
elif config_path.is_dir():
config_files.extend(config_path.glob('*.json'))
for cf in config_files:
with open(cf, 'r') as f:
model_cfg = json.load(f)
# 验证配置完整性
has_vision = all(k in model_cfg for k in ('embed_dim', 'vision_cfg', 'text_cfg'))
has_audio = all(k in model_cfg for k in ('embed_dim', 'audio_cfg', 'text_cfg'))
if has_vision or has_audio:
_MODEL_CONFIGS[cf.stem] = model_cfg
_rescan_model_configs() # 初始化时执行
```
### 模型名称解析 (第92-135行)
```python
def parse_model_name(model_name: str) -> Tuple[Optional[str], str]:
"""
解析模型名称，识别schema前缀
支持的schema：
- 'hf-hub:org/model' → HuggingFace Hub
- 'local-dir:/path' → 本地目录
- 'ViT-B-32' → 内置模型（无schema）
"""
# 1. 检查local-dir schema
if model_name.startswith(LOCAL_DIR_PREFIX):
identifier = model_name[len(LOCAL_DIR_PREFIX):]
if not identifier:
raise ValueError("Empty path after 'local-dir:'")
return 'local-dir', identifier
# 2. 检查hf-hub schema
elif model_name.startswith(HF_HUB_PREFIX):
identifier = model_name[len(HF_HUB_PREFIX):]
if not identifier:
raise ValueError("Empty identifier after 'hf-hub:'")
return 'hf-hub', identifier
# 3. 无schema → 内置模型名
else:
return None, model_name
```
### 配置获取函数 (第154-169行)
```python
def get_model_config(model_name):
"""
从指定位置获取模型配置
优先级：local-dir > hf-hub > 内置配置
"""
loc, model_id = parse_model_name(model_name)
if loc == 'local-dir':
local_path = Path(model_id) / 'open_clip_config.json'
with open(local_path, 'r', encoding='utf-8') as f:
config = json.load(f)
return config.get('model_cfg', config)
elif loc == 'hf-hub':
config = _get_hf_config(model_id) # 从HF下载
return config.get('model_cfg', config)
elif model_name in _MODEL_CONFIGS:
return deepcopy(_MODEL_CONFIGS[model_name])
else:
return None
```
### 权重加载 (第172-238行)
```python
def load_checkpoint(model, checkpoint_path, strict=True, device='cpu'):
"""
加载权重到模型，处理各种格式兼容性
"""
# 1. 根据扩展名选择加载方式
if str(checkpoint_path).endswith(".safetensors"):
from safetensors.torch import load_file
checkpoint = load_file(checkpoint_path, device=device)
else:
checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=True)
# 2. 提取state_dict
if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
state_dict = checkpoint['state_dict']
else:
state_dict = checkpoint
# 3. 处理DDP权重（移除'module.'前缀）
if next(iter(state_dict.items()))[0].startswith('module'):
state_dict = {k[7:]: v for k, v in state_dict.items()}
# 4. 第三方权重转换（如OpenAI、SigLIP）
state_dict = convert_state_dict(model, state_dict)
# 5. 处理旧格式兼容性
if 'positional_embedding' in state_dict and not hasattr(model, 'positional_embedding'):
state_dict = convert_to_custom_text_state_dict(state_dict)
# 6. 处理logit_scale形状差异
if 'logit_scale' in state_dict and model.logit_scale.ndim != state_dict['logit_scale'].ndim:
state_dict['logit_scale'] = state_dict['logit_scale'].reshape(model.logit_scale.shape)
# 7. 调整位置编码大小
resize_pos_embed(state_dict, model)
resize_text_pos_embed(state_dict, model)
# 8. 加载权重
incompatible_keys = model.load_state_dict(state_dict, strict=strict)
return incompatible_keys
```
### 主创建函数 (第264-702行)
```python
def create_model(
model_name: str,
pretrained: Optional[str] = None,
precision: str = 'fp32',
device: Union[str, torch.device] = 'cpu',
force_quick_gelu: bool = False,
force_custom_text: bool = False,
force_patch_dropout: Optional[float] = None,
force_image_size: Optional[Union[int, Tuple[int, int]]] = None,
pretrained_image: bool = False,
pretrained_text: bool = True,
pretrained_image_path: Optional[str] = None,
pretrained_text_path: Optional[str] = None,
cache_dir: Optional[str] = None,
require_pretrained: bool = False,
**model_kwargs,
) -> torch.nn.Module:
"""
创建CLIP模型的主工厂函数
参数说明：
- model_name: 模型标识符（可含schema前缀）
- pretrained: 预训练权重来源（仅在无schema时使用）
- force_quick_gelu: 强制使用QuickGELU
- force_custom_text: 强制使用自定义文本编码器
- pretrained_image/text: 是否加载基础塔权重
- pretrained_*_path: 单独加载指定塔的权重路径
- require_pretrained: 必须加载完整CLIP权重
"""
# 1. 解析模型名称
schema, identifier = parse_model_name(model_name)
# 2. 根据schema获取配置
model_cfg = None
checkpoint_path = None
preprocess_cfg = asdict(PreprocessCfg())
if schema == 'local-dir':
# 从本地目录加载
model_cfg = _load_local_config(identifier)
checkpoint_path = _find_checkpoint_in_dir(Path(identifier))
elif schema == 'hf-hub':
# 从HuggingFace加载
hf_config = _get_hf_config(identifier)
model_cfg = hf_config['model_cfg']
checkpoint_path = download_pretrained_from_hf(identifier)
else: # 无schema → 内置模型
model_cfg = get_model_config(identifier)
if pretrained:
checkpoint_path = download_pretrained(pretrained)
# 3. 应用配置覆盖
if force_quick_gelu:
model_cfg["quick_gelu"] = True
if force_patch_dropout is not None:
model_cfg['vision_cfg']["patch_dropout"] = force_patch_dropout
if force_image_size is not None:
model_cfg['vision_cfg']["image_size"] = force_image_size
# 4. 确定模型类型
if 'audio_cfg' in model_cfg:
model_class = CLAP
elif "multimodal_cfg" in model_cfg:
model_class = CoCa
elif force_custom_text or 'hf_model_name' in model_cfg.get('text_cfg', {}):
model_class = CustomTextCLIP
else:
model_class = CLIP
# 5. 创建模型实例
model = model_class(**model_cfg)
# 6. 设置设备和精度
_set_model_device_and_precision(model, device, precision)
# 7. 加载权重
if checkpoint_path and load_weights:
load_checkpoint(model, checkpoint_path)
# 8. 加载塔特定权重（可选）
if pretrained_image_path:
image_state_dict = load_state_dict(pretrained_image_path)
model.visual.load_state_dict(image_state_dict, strict=False)
if pretrained_text_path:
text_state_dict = load_state_dict(pretrained_text_path)
model.text.load_state_dict(text_state_dict, strict=False)
# 9. 设置预处理配置
set_model_preprocess_cfg(model, preprocess_cfg)
return model
```
### 分词器获取 (第746-886行)
```python
def get_tokenizer(
model_name: str = '',
context_length: Optional[int] = None,
**kwargs,
):
"""
根据模型配置获取对应的分词器
分词器类型选择逻辑：
1. tiktoken → TikTokenTokenizer (GenLIP)
2. hf_tokenizer_name → HFTokenizer
3. siglip命名 → SigLipTokenizer
4. 默认 → SimpleTokenizer
"""
# 1. 获取配置
config = get_model_config(model_name) or {}
text_config = config.get('text_cfg', {})
# 2. 确定context_length
if context_length is None:
context_length = text_config.get('context_length', DEFAULT_CONTEXT_LENGTH)
# 3. 选择分词器类型
tokenizer_type = text_config.get('tokenizer_type', '')
hf_tokenizer_name = text_config.get('hf_tokenizer_name', '')
if tokenizer_type == 'tiktoken':
return TikTokenTokenizer(
encoding_name=text_config.get('tiktoken_name', 'cl100k_base'),
context_length=context_length,
**kwargs
)
elif hf_tokenizer_name:
return HFTokenizer(
tokenizer_source=hf_tokenizer_name,
context_length=context_length,
**kwargs
)
elif 'siglip' in model_name.lower():
return SigLipTokenizer(context_length=context_length)
else:
return SimpleTokenizer(context_length=context_length, **kwargs)
```
### 创建模型与变换 (第1081-1214行)
```python
def create_model_and_transforms(
model_name: str,
pretrained: Optional[str] = None,
precision: str = 'fp32',
device: Union[str, torch.device] = 'cpu',
aug_cfg: Optional[Union[Dict, AugmentationCfg]] = None,
audio_aug_cfg: Optional[Any] = None,
**kwargs,
):
"""
同时创建模型和预处理变换
返回: (model, preprocess_train, preprocess_val)
"""
# 1. 创建模型
model = create_model(
model_name,
pretrained,
precision=precision,
device=device,
**kwargs,
)
# 2. 构建预处理变换
preprocess_train, preprocess_val = _build_preprocess(
model,
aug_cfg=aug_cfg,
audio_aug_cfg=audio_aug_cfg
)
return model, preprocess_train, preprocess_val
```
|
|
## model.py 详解
### 配置数据类 (第37-151行)
```python
@dataclass
class CLIPVisionCfg:
"""视觉编码器配置"""
layers: Union[Tuple[int, int, int, int], int] = 12 # ResNet层或ViT层数
width: int = 768 # ViT宽度
head_width: int = 64 # 注意力头宽度
mlp_ratio: float = 4.0 # MLP扩展比
patch_size: int = 16 # ViT patch大小
image_size: Union[Tuple[int, int], int] = 224 # 输入图像大小
patch_dropout: float = 0. # patch丢弃率
attentional_pool: bool = False # 是否使用注意力池化
pos_embed_type: str = 'learnable' # 位置编码类型
pool_type: str = 'tok' # 池化类型
# timm相关配置
timm_model_name: Optional[str] = None # timm模型名
timm_model_pretrained: bool = False # 是否使用预训练
timm_pool: str = 'avg' # timm池化方式
@dataclass
class CLIPTextCfg:
"""文本编码器配置"""
text_arch: str = "clip" # 架构类型
context_length: int = 77 # 上下文长度
vocab_size: int = 49408 # 词表大小
hf_tokenizer_name: Optional[str] = None # HF分词器名
tokenizer_type: str = '' # 分词器类型
width: int = 512 # Transformer宽度
heads: int = 8 # 注意力头数
layers: int = 12 # Transformer层数
mlp_ratio: float = 4.0 # MLP扩展比
pool_type: str = 'argmax' # 池化类型
proj_type: str = 'linear' # 投影类型
# HuggingFace相关
hf_model_name: Optional[str] = None # HF模型名
hf_model_pretrained: bool = True # 是否使用预训练
hf_pooler_type: str = 'mean_pooler' # HF池化器
```
### 视觉塔构建 (第171-248行)
```python
def _build_vision_tower(
embed_dim: int,
vision_cfg: CLIPVisionCfg,
quick_gelu: bool = False,
cast_dtype: Optional[torch.dtype] = None
):
"""构建视觉编码器塔"""
if isinstance(vision_cfg, dict):
vision_cfg = CLIPVisionCfg(**vision_cfg)
act_layer = QuickGELU if quick_gelu else nn.GELU
# 分支1: timm模型
if vision_cfg.timm_model_name:
visual = TimmModel(
vision_cfg.timm_model_name,
pretrained=vision_cfg.timm_model_pretrained,
pool=vision_cfg.timm_pool,
proj=vision_cfg.timm_proj,
embed_dim=embed_dim,
image_size=vision_cfg.image_size,
)
# 分支2: ResNet
elif isinstance(vision_cfg.layers, (tuple, list)):
vision_heads = vision_cfg.width * 32 // vision_cfg.head_width
visual = ModifiedResNet(
layers=vision_cfg.layers,
output_dim=embed_dim,
heads=vision_heads,
image_size=vision_cfg.image_size,
width=vision_cfg.width,
)
# 分支3: Vision Transformer
else:
vision_heads = vision_cfg.width // vision_cfg.head_width
norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
visual = VisionTransformer(
image_size=vision_cfg.image_size,
patch_size=vision_cfg.patch_size,
width=vision_cfg.width,
layers=vision_cfg.layers,
heads=vision_heads,
mlp_ratio=vision_cfg.mlp_ratio,
patch_dropout=vision_cfg.patch_dropout,
attentional_pool=vision_cfg.attentional_pool,
pos_embed_type=vision_cfg.pos_embed_type,
pool_type=vision_cfg.pool_type,
output_dim=embed_dim,
act_layer=act_layer,
norm_layer=norm_layer,
)
return visual
```
### 文本塔构建 (第251-315行)
```python
def _build_text_tower(
embed_dim: int,
text_cfg: CLIPTextCfg,
quick_gelu: bool = False,
cast_dtype: Optional[torch.dtype] = None,
):
"""构建文本编码器塔"""
if isinstance(text_cfg, dict):
text_cfg = CLIPTextCfg(**text_cfg)
# 分支1: ModernTextTransformer
if text_cfg.text_arch == "modern":
text = ModernTextTransformer(text_cfg, output_dim=embed_dim)
# 分支2: HuggingFace编码器
elif text_cfg.hf_model_name:
text = HFTextEncoder(
text_cfg.hf_model_name,
output_dim=embed_dim,
proj_type=text_cfg.hf_proj_type,
pooler_type=text_cfg.hf_pooler_type,
pretrained=text_cfg.hf_model_pretrained,
)
# 分支3: 原生CLIP Transformer
else:
act_layer = QuickGELU if quick_gelu else nn.GELU
norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
text = TextTransformer(
context_length=text_cfg.context_length,
vocab_size=text_cfg.vocab_size,
width=text_cfg.width,
heads=text_cfg.heads,
layers=text_cfg.layers,
mlp_ratio=text_cfg.mlp_ratio,
output_dim=embed_dim,
pool_type=text_cfg.pool_type,
proj_type=text_cfg.proj_type,
act_layer=act_layer,
norm_layer=norm_layer,
)
text.variable_text = bool(text_cfg.variable_text)
return text
```
### CLIP主类 (第318-548行)
```python
class CLIP(nn.Module):
"""
CLIP主模型类 - 使用原生CLIP文本塔
结构：
- visual: 视觉编码器（ResNet/ViT/timm）
- transformer: 文本Transformer
- logit_scale: 可学习温度参数
"""
def __init__(
self,
embed_dim: int,
vision_cfg: CLIPVisionCfg,
text_cfg: CLIPTextCfg,
quick_gelu: bool = False,
init_logit_scale: float = np.log(1 / 0.07), # ≈ 2.659
init_logit_bias: Optional[float] = None,
cast_dtype: Optional[torch.dtype] = None,
output_dict: bool = False,
):
super().__init__()
self.output_dict = output_dict
self.embed_dim = embed_dim
# 构建视觉塔
self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
# 构建文本塔（内联到CLIP中）
text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
self.transformer = text.transformer
self.context_length = text.context_length
self.vocab_size = text.vocab_size
self.token_embedding = text.token_embedding
self.positional_embedding = text.positional_embedding
self.ln_final = text.ln_final
self.text_projection = text.text_projection
self.text_pool_type = text.pool_type
self.attn_mask = text.attn_mask
# 可学习参数
self.logit_scale = nn.Parameter(torch.ones([]) * init_logit_scale)
if init_logit_bias is not None:
self.logit_bias = nn.Parameter(torch.ones([]) * init_logit_bias)
else:
self.logit_bias = None
def _encode_image(self, image, normalize: bool = False):
"""编码图像"""
features = self.visual(image)
return F.normalize(features, dim=-1) if normalize else features
def _encode_text(self, text, normalize: bool = False):
"""编码文本"""
cast_dtype = self.transformer.get_cast_dtype()
# Token嵌入
x = self.token_embedding(text).to(cast_dtype)
# 加位置编码
x = x + self.positional_embedding.to(cast_dtype)
# Transformer编码
x = self.transformer(x, attn_mask=self.attn_mask)
# 最终LayerNorm
x = self.ln_final(x)
# 全局池化
x = text_global_pool(x, text, self.text_pool_type)
# 投影到共享空间
if self.text_projection is not None:
if isinstance(self.text_projection, nn.Linear):
x = self.text_projection(x)
else:
x = x @ self.text_projection
return F.normalize(x, dim=-1) if normalize else x
def get_logits(self, image, text):
"""计算图像-文本logits"""
image_features = self._encode_image(image, normalize=True)
text_features = self._encode_text(text, normalize=True)
# 余弦相似度 × 温度
image_logits = self.logit_scale.exp() * (image_features @ text_features.T)
if self.logit_bias is not None:
image_logits += self.logit_bias
text_logits = image_logits.T
return image_logits, text_logits
def forward(self, image=None, text=None):
"""前向传播"""
image_features = self._encode_image(image, normalize=True) if image is not None else None
text_features = self._encode_text(text, normalize=True) if text is not None else None
if self.output_dict:
return {
"image_features": image_features,
"text_features": text_features,
"logit_scale": self.logit_scale.exp()
}
if self.logit_bias is not None:
return image_features, text_features, self.logit_scale.exp(), self.logit_bias.clone()
return image_features, text_features, self.logit_scale.exp()
```
### CustomTextCLIP类 (第551-735行)
```python
class CustomTextCLIP(nn.Module):
"""
使用独立文本塔的CLIP变体
与CLIP的区别：
- 文本塔作为self.text独立存在
- 支持HuggingFace预训练模型
- 支持可变长度文本
"""
def __init__(
self,
embed_dim: int,
vision_cfg: CLIPVisionCfg,
text_cfg: CLIPTextCfg,
quick_gelu: bool = False,
init_logit_scale: float = np.log(1 / 0.07),
**kwargs,
):
super().__init__()
self.embed_dim = embed_dim
# 构建双塔
self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu)
self.text = _build_text_tower(embed_dim, text_cfg, quick_gelu)
self.context_length = self.text.context_length
self.vocab_size = self.text.vocab_size
# 可学习参数
self.logit_scale = nn.Parameter(torch.ones([]) * init_logit_scale)
self.logit_bias = nn.Parameter(torch.ones([]) * init_logit_bias) if init_logit_bias else None
def _encode_text(self, text, normalize: bool = False):
"""编码文本（委托给self.text）"""
features = self.text(text)
return F.normalize(features, dim=-1) if normalize else features
```
### 权重转换函数 (第738-767行)
```python
def convert_weights_to_lp(model: nn.Module, dtype=torch.float16):
"""将模型权重转换为低精度（fp16/bf16）"""
def _convert_weights(l):
# 卷积和线性层
if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
l.weight.data = l.weight.data.to(dtype)
if l.bias is not None:
l.bias.data = l.bias.data.to(dtype)
# 注意力层
if isinstance(l, Attention):
for attr in ['in_proj_weight', 'q_proj_weight', 'k_proj_weight', 'v_proj_weight']:
tensor = getattr(l, attr, None)
if tensor is not None:
tensor.data = tensor.data.to(dtype)
# 文本投影
if isinstance(l, (CLIP, TextTransformer)):
attr = getattr(l, "text_projection", None)
if attr is not None:
attr.data = attr.data.to(dtype)
# 视觉投影
if isinstance(l, VisionTransformer):
attr = getattr(l, "proj", None)
if attr is not None:
attr.data = attr.data.to(dtype)
model.apply(_convert_weights)
```
### 位置编码调整 (第790-854行)
```python
def resize_pos_embed(state_dict, model, interpolation='bicubic', antialias=True):
"""
调整位置编码大小
当加载不同分辨率的权重时使用
"""
old_pos_embed = state_dict.get('visual.positional_embedding', None)
if old_pos_embed is None or not hasattr(model.visual, 'grid_size'):
return
grid_size = to_2tuple(model.visual.grid_size)
extra_tokens = 1 # CLS token
new_seq_len = grid_size[0] * grid_size[1] + extra_tokens
if new_seq_len == old_pos_embed.shape[0]:
return # 无需调整
# 分离CLS token和位置编码
pos_emb_tok, pos_emb_img = old_pos_embed[:extra_tokens], old_pos_embed[extra_tokens:]
# 旧grid大小
old_grid_size = to_2tuple(int(math.sqrt(len(pos_emb_img))))
# 双三次插值调整
pos_emb_img = pos_emb_img.reshape(1, old_grid_size[0], old_grid_size[1], -1).permute(0, 3, 1, 2)
pos_emb_img = F.interpolate(
pos_emb_img,
size=grid_size,
mode=interpolation,
antialias=antialias,
)
pos_emb_img = pos_emb_img.permute(0, 2, 3, 1).reshape(1, grid_size[0] * grid_size[1], -1)[0]
# 拼接CLS token
new_pos_embed = torch.cat([pos_emb_tok, pos_emb_img], dim=0)
state_dict['visual.positional_embedding'] = new_pos_embed
```
|
|
## loss.py 详解
### 分布式特征收集 (第29-54行)
```python
def gather_features(
image_features,
text_features,
local_loss=False,
gather_with_grad=False,
rank=0,
world_size=1,
):
"""
在多GPU训练中收集所有进程的特征
行为：
- gather_with_grad=True: 保留梯度收集（内存开销大）
- gather_with_grad=False: 收集后为本地rank添加梯度
"""
if gather_with_grad:
all_image_features = _all_gather_with_grad(image_features)
all_text_features = _all_gather_with_grad(text_features)
else:
# 创建收集容器
gathered_image_features = [torch.zeros_like(image_features) for _ in range(world_size)]
gathered_text_features = [torch.zeros_like(text_features) for _ in range(world_size)]
# all_gather收集
dist.all_gather(gathered_image_features, image_features)
dist.all_gather(gathered_text_features, text_features)
if not local_loss:
# 确保本地rank有梯度
gathered_image_features[rank] = image_features
gathered_text_features[rank] = text_features
all_image_features = torch.cat(gathered_image_features, dim=0)
all_text_features = torch.cat(gathered_text_features, dim=0)
return all_image_features, all_text_features
```
### ClipLoss类 (第57-141行)
```python
class ClipLoss(nn.Module):
"""
标准CLIP对比损失
损失 = (CrossEntropy(image_logits, labels) + CrossEntropy(text_logits, labels)) / 2
"""
def __init__(
self,
local_loss=False, # 是否使用局部损失
gather_with_grad=False, # 收集时是否保留梯度
cache_labels=False, # 是否缓存标签
rank=0, # 进程rank
world_size=1, # 总进程数
):
super().__init__()
self.local_loss = local_loss
self.gather_with_grad = gather_with_grad
self.cache_labels = cache_labels
self.rank = rank
self.world_size = world_size
def get_ground_truth(self, device, num_logits) -> torch.Tensor:
"""获取真实标签（对角线）"""
if self.prev_num_logits != num_logits or device not in self.labels:
labels = torch.arange(num_logits, device=device, dtype=torch.long)
if self.world_size > 1 and self.local_loss:
labels = labels + num_logits * self.rank
if self.cache_labels:
self.labels[device] = labels
self.prev_num_logits = num_logits
else:
labels = self.labels[device]
return labels
def get_logits(self, image_features, text_features, logit_scale, logit_bias=None):
"""计算logits"""
if self.world_size > 1:
all_image_features, all_text_features = gather_features(...)
if self.local_loss:
# 局部损失：本地图像 × 全局文本
logits_per_image = logit_scale * image_features @ all_text_features.T
logits_per_text = logit_scale * text_features @ all_image_features.T
else:
# 全局损失：全局图像 × 全局文本
logits_per_image = logit_scale * all_image_features @ all_text_features.T
logits_per_text = logits_per_image.T
else:
# 单GPU
logits_per_image = logit_scale * image_features @ text_features.T
logits_per_text = logit_scale * text_features @ image_features.T
if logit_bias is not None:
logits_per_image += logit_bias
logits_per_text += logit_bias
return logits_per_image, logits_per_text
def forward(self, image_features, text_features, logit_scale, logit_bias=None):
"""
计算CLIP损失
返回: 总损失 = (image_to_text_loss + text_to_image_loss) / 2
"""
# 1. 计算logits
logits_per_image, logits_per_text = self.get_logits(
image_features, text_features, logit_scale, logit_bias
)
# 2. 获取标签（对角线）
labels = self.get_ground_truth(image_features.device, logits_per_image.shape[0])
# 3. 计算双向交叉熵
total_loss = (
F.cross_entropy(logits_per_image, labels) + # 图像→文本
F.cross_entropy(logits_per_text, labels) # 文本→图像
) / 2
return total_loss
```
**损失计算示意图：**
```
相似度矩阵 (batch_size=4):
文本1 文本2 文本3 文本4
图像1 [ 9.2 -3.1 -4.5 -5.2 ] ← 匹配
图像2 [ -2.8 8.7 -3.9 -4.1 ] ← 匹配
图像3 [ -4.2 -3.5 9.1 -3.8 ] ← 匹配
图像4 [ -5.1 -4.8 -3.7 8.9 ] ← 匹配
标签: [0, 1, 2, 3] (对角线索引)
图像→文本损失: -log(exp(9.2) / (exp(9.2)+exp(-3.1)+exp(-4.5)+exp(-5.2)))
文本→图像损失: -log(exp(9.2) / (exp(9.2)+exp(-2.8)+exp(-4.2)+exp(-5.1)))
总损失 = (图像损失 + 文本损失) / 2
```
### CoCaLoss类 (第144-184行)
```python
class CoCaLoss(ClipLoss):
"""
CoCa (Contrastive Captioners) 损失
组合损失 = λ₁ × ClipLoss + λ₂ × CaptionLoss
"""
def __init__(
self,
caption_loss_weight,
clip_loss_weight,
pad_id=0,
local_loss=False,
gather_with_grad=False,
**kwargs,
):
super().__init__(local_loss, gather_with_grad, **kwargs)
self.clip_loss_weight = clip_loss_weight
self.caption_loss_weight = caption_loss_weight
self.caption_loss = nn.CrossEntropyLoss(ignore_index=pad_id)
def forward(self, image_features, text_features, logits, labels, logit_scale):
"""
参数:
- logits: 生成式logits [batch, seq_len, vocab_size]
- labels: 文本标签 [batch, seq_len]
"""
# 1. 对比损失
if self.clip_loss_weight:
clip_loss = super().forward(image_features, text_features, logit_scale)
clip_loss = self.clip_loss_weight * clip_loss
# 2. 生成式损失（next token prediction）
caption_loss = self.caption_loss(
logits.permute(0, 2, 1), # [batch, vocab, seq_len]
labels,
)
caption_loss = self.caption_loss_weight * caption_loss
return clip_loss, caption_loss
```
### SigLipLoss类 (第314-489行)
```python
class SigLipLoss(nn.Module):
"""
Sigmoid Loss for Language Image Pre-training
与CLIP损失的区别：
- 使用Sigmoid代替Softmax
- 不需要全局归一化
- 更适合大规模分布式训练
"""
def __init__(
self,
cache_labels: bool = False,
rank: int = 0,
world_size: int = 1,
dist_impl: Optional[str] = None, # 'bidir', 'shift', 'reduce', 'gather'
chunk_size: int = 0,
):
super().__init__()
self.rank = rank
self.world_size = world_size
self.dist_impl = dist_impl or 'bidir'
self.chunk_size = chunk_size
def get_ground_truth(self, device, dtype, num_logits, negative_only=False):
"""生成标签矩阵 (+1对角线, -1其他)"""
labels = -torch.ones((num_logits, num_logits), device=device, dtype=dtype)
if not negative_only:
labels = 2 * torch.eye(num_logits, device=device, dtype=dtype) + labels
return labels
def _loss(self, image_features, text_features, logit_scale, logit_bias=None, negative_only=False):
"""计算sigmoid loss"""
logits = logit_scale * image_features @ text_features.T
if logit_bias is not None:
logits += logit_bias
labels = self.get_ground_truth(
image_features.device,
image_features.dtype,
image_features.shape[0],
negative_only=negative_only,
)
# -logsigmoid(labels * logits) 当labels=+1时为正样本，=-1时为负样本
loss = -F.logsigmoid(labels * logits).sum() / image_features.shape[0]
return loss
def forward(self, image_features, text_features, logit_scale, logit_bias):
"""前向计算（含分布式）"""
loss = self._loss(image_features, text_features, logit_scale, logit_bias)
# 分布式邻居交换
if self.world_size > 1:
if self.dist_impl == 'bidir':
# 双向交换
right_rank = (self.rank + 1) % self.world_size
left_rank = (self.rank - 1 + self.world_size) % self.world_size
for i in range((self.world_size - 1) // 2):
text_features_recv = neighbour_exchange_bidir_with_grad(...)
for f in text_features_recv:
loss += self._loss(image_features, f, logit_scale, negative_only=True)
elif self.dist_impl == 'shift':
# 环形移位
for i in range(self.world_size - 1):
text_features_from_left = neighbour_exchange_with_grad(...)
loss += self._loss(image_features, text_features_from_left, ...)
return loss
```
### GenLipLoss类 (第542-561行)
```python
class GenLipLoss(nn.Module):
"""
生成式CLIP损失（用于GenLIP模型）
纯自回归语言建模损失
"""
def __init__(self, ignore_index: int = -100):
super().__init__()
self.ignore_index = ignore_index
self.caption_loss = nn.CrossEntropyLoss(ignore_index=ignore_index)
def forward(self, logits, labels):
"""
参数:
- logits: [batch, seq_len, vocab_size]
- labels: [batch, seq_len]
损失在展平的token上计算
"""
loss = self.caption_loss(
logits.reshape(-1, logits.shape[-1]), # [batch*seq_len, vocab_size]
labels.reshape(-1), # [batch*seq_len]
)
return loss
```
|
|
## 训练全流程
### 训练入口 (main.py)
```python
def main(args):
"""OpenCLIP训练主函数"""
# 1. 初始化分布式环境
device = init_distributed_device(args)
# 2. 设置实验名称和日志
args.name = f"{date_str}-model_{model_name}-lr_{lr}-b_{batch_size}"
setup_logging(args.log_path, args.log_level)
# 3. 创建模型
model, preprocess_train, preprocess_val = create_model_and_transforms(
args.model,
args.pretrained,
precision=args.precision,
device=device,
aug_cfg=args.aug_cfg,
)
# 4. 创建训练任务（封装模型+损失）
task = create_task(args, model=model)
# 5. 创建数据加载器
data = get_data(args, preprocess_train, preprocess_val)
# 6. 创建优化器
optimizer = create_optimizer(args, model)
# 7. 创建学习率调度器
scheduler = cosine_lr(optimizer, args.lr, args.warmup, args.epochs)
# 8. 创建训练状态
state = TrainState(
task=task,
optimizer=optimizer,
scheduler=scheduler,
scaler=scaler,
epoch=start_epoch,
)
# 9. 训练循环
for epoch in range(start_epoch, args.epochs):
train_one_epoch(state, data, args)
evaluate(task, data, epoch, args)
# 10. 保存检查点
save_checkpoint(state, checkpoint_path)
```
### 训练状态 (train.py)
```python
@dataclass
class TrainState:
"""训练运行时状态"""
task: Any # 训练任务
optimizer: Optional[torch.optim.Optimizer] = None
scaler: Optional[Any] = None # AMP GradScaler
scheduler: Optional[Callable] = None # 学习率调度器
epoch: int = 0 # 当前epoch
global_step: int = 0 # 全局步数
samples_seen: int = 0 # 已处理样本数
compiled_train_step: Optional[Callable] = None # 编译后的训练步
losses_ema: dict = field(default_factory=dict) # 损失EMA
```
### 单epoch训练 (train.py 第337-520行)
```python
def train_one_epoch(state: TrainState, data, args, tb_writer=None):
"""训练一个epoch"""
task = state.task
optimizer = state.optimizer
scheduler = state.scheduler
epoch = state.epoch
# 1. 设置模型为训练模式
task.train()
# 2. 初始化数据加载器
dataloader = data['train'].dataloader
data['train'].set_epoch(epoch)
# 3. 确定训练步函数
if args.torchcompile and args.torchcompile_strategy == 'step':
train_step = _get_compiled_train_step(state, autocast, args)
elif args.accum_freq > 1:
train_step = None # 使用累积
else:
train_step = _make_train_step_no_accum_no_scaler(...)
# 4. 训练循环
for i, batch in enumerate(dataloader):
# 4.1 更新学习率
step = num_batches_per_epoch * epoch + i // args.accum_freq
scheduler(step)
# 4.2 准备批次
batch = task.prepare_batch(batch, device=device, input_dtype=input_dtype)
# 4.3 训练步
if train_step is not None:
optimizer.zero_grad()
losses, report = train_step(batch)
task.clamp_logit_scale()
else:
# 使用梯度累积或混合精度
result = _train_step_eager(task, batch, accum_state, optimizer, scaler, autocast, args)
losses, report = result
# 4.4 更新计数器
state.global_step = step + 1
state.samples_seen += batch_size * args.world_size
# 4.5 日志记录
if is_master(args) and i % args.log_every_n_steps == 0:
_logger.info(
f"Epoch: {epoch} "
f"Loss: {losses['loss'].item():.4f} "
f"LR: {get_learning_rate(optimizer):.5f} "
f"Logit Scale: {report['logit_scale'].item():.3f}"
)
# TensorBoard记录
if tb_writer is not None:
for name, val in losses.items():
tb_writer.add_scalar(f"train/{name}", val.item(), step)
# 5. Epoch结束日志
_logger.info(f"End epoch {epoch} | Avg Loss: {loss_avg:.4f}")
```
### 训练步函数 (train.py 第163-216行)
```python
def _make_train_step_no_accum_no_scaler(task, optimizer, autocast, args):
"""
构建单步训练函数（无累积、无scaler）
这是torch.compile的目标函数
"""
grad_clip_norm = getattr(args, "grad_clip_norm", None)
clip_params = tuple(task.trainable_module.parameters()) if grad_clip_norm else ()
def train_step(batch):
# 1. 获取损失缩放（NaFlex）
loss_scale = get_naflex_loss_scale(batch, args, task)
# 2. 前向传播
with autocast():
losses, report = task(batch)
total_loss = losses["loss"]
if loss_scale != 1.0:
total_loss = total_loss * loss_scale
# 3. 反向传播
total_loss.backward()
# 4. 梯度裁剪
if grad_clip_norm is not None:
torch.nn.utils.clip_grad_norm_(clip_params, grad_clip_norm)
# 5. 优化器步
optimizer.step()
return losses, report
return train_step
```
### 急切模式训练步 (train.py 第221-311行)
```python
def _train_step_eager(task, batch, accum_state, optimizer, scaler, autocast, args):
"""
急切模式训练步（支持梯度累积和AMP）
返回: (losses, report, batch_size, accum_state)
"""
# === 无梯度累积 ===
if args.accum_freq == 1:
optimizer.zero_grad()
with autocast():
losses, report = task(batch)
total_loss = losses["loss"]
loss_scale = get_naflex_loss_scale(batch, args, task)
if loss_scale != 1.0:
total_loss = total_loss * loss_scale
backward(total_loss, scaler)
_finish_eager_train_step(task, optimizer, scaler, args)
return losses, report, batch_size, accum_state
# === 有梯度累积 ===
accum_batches, accum_features = accum_state
# 1. 缓存特征（无梯度）
with torch.no_grad():
with autocast():
model_out = task.trainable_module(**batch)
model_out.pop("logit_scale", None)
model_out.pop("logit_bias", None)
for key, val in model_out.items():
if key in accum_features:
accum_features[key].append(val)
else:
accum_features[key] = [val]
accum_batches.append(batch)
# 2. 等待累积完成
if len(accum_batches) < args.accum_freq:
return None
# 3. 计算累积梯度
optimizer.zero_grad()
for j in range(args.accum_freq):
batch_j = accum_batches[j]
is_last_step = (j == args.accum_freq - 1)
# 禁用梯度同步（FSDP/DDP）
use_no_sync = not is_last_step and (isinstance(task.trainable_module, DistributedDataParallel))
ddp_context = task.trainable_module.no_sync() if use_no_sync else nullcontext()
with ddp_context:
with autocast():
model_out = task.trainable_module(**batch_j)
# 重新组装输入（使用累积特征）
inputs_no_accum = {"logit_scale": model_out.pop("logit_scale")}
inputs = {}
for key, val in accum_features.items():
accumulated = accum_features[key]
inputs[key] = torch.cat(accumulated[:j] + [model_out[key]] + accumulated[j + 1:])
losses, report = task.compute_accum_loss(inputs, inputs_no_accum, accum_batches)
total_loss = sum(v for k, v in losses.items() if k.endswith('_loss'))
backward(total_loss, scaler)
# 4. 优化器步
_finish_eager_train_step(task, optimizer, scaler, args)
return losses, report, step_batch_size, ([], {})
```
### 学习率调度 (scheduler.py)
```python
def cosine_lr(optimizer, base_lr, warmup_lr, warmup_steps, total_steps):
"""
余弦退火学习率调度器
公式:
- warmup阶段: lr = warmup_lr + (base_lr - warmup_lr) * step / warmup_steps
- decay阶段: lr = 0.5 * (1 + cos(π * (step - warmup_steps) / (total_steps - warmup_steps)))
"""
def step_fn(step):
if step < warmup_steps:
# 线性warmup
lr = warmup_lr + (base_lr - warmup_lr) * step / warmup_steps
else:
# 余弦衰减
lr = 0.5 * base_lr * (1 + math.cos(math.pi * (step - warmup_steps) / (total_steps - warmup_steps)))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
return lr
return step_fn
```
### 优化器创建 (optim.py)
```python
def create_optimizer(args, model):
"""
创建优化器
支持:
- AdamW (默认)
- SGD
- AdaW (带层衰减)
"""
# 跳过不需要衰减的参数
skip = {}
if hasattr(model, 'no_weight_decay'):
for n, p in model.named_parameters():
if n in model.no_weight_decay():
skip[p] = True
# 参数分组（不同学习率）
decay_params = []
no_decay_params = []
for n, p in model.named_parameters():
if p not in skip:
if 'bias' in n or 'norm' in n or 'ln_' in n:
no_decay_params.append(p)
else:
decay_params.append(p)
optimizer_args = dict(
params=[
dict(params=decay_params, weight_decay=args.wd),
dict(params=no_decay_params, weight_decay=0)
],
lr=args.lr,
betas=(args.beta1, args.beta2),
eps=args.eps,
)
if args.optim == 'adamw':
return torch.optim.AdamW(**optimizer_args)
elif args.optim == 'sgd':
return torch.optim.SGD(**optimizer_args)
else:
raise ValueError(f"Unknown optimizer: {args.optim}")
```
### 训练命令示例
```bash
# 基础训练
python -m open_clip_train.main \
--train-data "pipe:curl -s -L {url} | tar xf -" \
--train-num-samples 1000000 \
--dataset-type webdataset \
--batch-size 256 \
--epochs 32 \
--lr 1e-3 \
--warmup 2000 \
--model ViT-B-32 \
--name my_experiment
# 带验证的训练
python -m open_clip_train.main \
--train-data "/path/to/train shard-{0000..0999}.tar" \
--val-data "/path/to/val" \
--dataset-type webdataset \
--imagenet-val "/path/to/imagenet/val" \
--batch-size 128 \
--workers 8 \
--epochs 64 \
--lr 5e-4 \
--model ViT-L-14 \
--pretrained laion2b_s34b_b79k \
--report-to wandb tensorboard
# FSDP2分布式训练
python -m open_clip_train.main \
--train-data "pipe:aws s3 cp ..." \
--batch-size 2048 \
--accum-freq 1 \
--fsdp \
--fsdp-shard-strategy FULL_SHARD \
--model ViT-H-14 \
--precision amp_bf16 \
--torchcompile-strategy step \
--grad-checkpointing
```
### 训练数据加载 (data.py)
```python
def get_data(args, preprocess_train, preprocess_val):
"""
创建训练和验证数据加载器
支持的数据源:
- WebDataset (.tar文件)
- CSV文件
- ImageFolder
- TFDS (部分)
"""
data = {}
# === 训练数据 ===
if args.dataset_type == 'webdataset':
# WebDataset管道
pipeline = [
wds.SimpleShardList(args.train_data),
wds.split_by_node,
wds.split_by_worker,
wds.tarfile_to_samples,
wds.decode("pil"),
wds.rename(image="jpg;png;jpeg;webp", text="txt"),
wds.map(filter_no_caption_or_no_image),
wds.map_dict(text=TokenizeText(tokenizer)),
wds.map_dict(image=preprocess_train, text=None),
wds.batched(args.batch_size, partial=False),
]
dataloader = wds.WebLoader(pipeline, num_workers=args.workers)
elif args.dataset_type == 'csv':
# CSV数据集
dataset = CsvDataset(
args.train_data,
transforms=preprocess_train,
img_key=args.csv_img_key,
caption_key=args.csv_caption_key,
tokenizer=tokenizer,
)
sampler = DistributedSampler(dataset, shuffle=True)
dataloader = DataLoader(dataset, batch_size=args.batch_size, sampler=sampler)
data['train'] = DataInfo(dataloader=dataloader, sampler=sampler)
# === 验证数据 ===
if args.val_dataset:
val_dataset = ImageFolder(args.val_dataset, transform=preprocess_val)
val_loader = DataLoader(val_dataset, batch_size=args.batch_size)
data['val'] = DataInfo(dataloader=val_loader)
# === ImageNet验证 (零样本) ===
if args.imagenet_val:
data['imagenet-val'] = get_imagenet(args, (preprocess_train, preprocess_val), 'val')
return data
```
|
|
## 推理与评估
### 基础推理示例
```python
import torch
import open_clip
from PIL import Image
# 1. 加载模型
model, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
model.eval()
tokenizer = open_clip.get_tokenizer('ViT-B-32')
# 2. 准备输入
image = preprocess(Image.open("image.jpg")).unsqueeze(0)
text = tokenizer(["a dog", "a cat", "a car"])
# 3. 推理
with torch.no_grad(), torch.cuda.amp.autocast():
image_features = model.encode_image(image)
text_features = model.encode_text(text)
# 4. 计算相似度
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)
```
### 零样本分类 (zero_shot.py)
```python
def zero_shot_eval(task, data, epoch, args, tokenizer=None):
"""ImageNet零样本评估"""
if "imagenet-val" not in data:
return {}
# 1. 获取ImageNet类别名
classnames = open_clip.IMAGENET_CLASSNAMES
templates = open_cipher.OPENAI_IMAGENET_TEMPLATES
# 2. 构建分类器
classifier = build_zero_shot_classifier(
model,
tokenizer=tokenizer,
classnames=classnames,
templates=templates,
device=args.device,
)
# 3. 评估
metrics = {}
dataloader = data['imagenet-val'].dataloader
with torch.inference_mode():
top1, top5 = 0.0, 0.0
for images, labels in dataloader:
images = images.to(args.device)
labels = labels.to(args.device)
# 前向传播
with torch.cuda.amp.autocast():
logits = 100.0 * model.logit_scale.exp() * (model.encode_image(images) @ classifier.T)
# 计算准确率
acc1, acc5 = accuracy(logits, labels, topk=(1, 5))
top1 += acc1
top5 += acc5
metrics['imagenet_val_top1'] = top1 / len(dataloader)
metrics['imagenet_val_top5'] = top5 / len(dataloader)
return metrics
```
### 检索评估 (metrics.py)
```python
def get_clip_metrics(
image_features,
text_features,
logit_scale,
image_key="image",
text_key="text",
retrieval_chunk_size=DEFAULT_RETRIEVAL_CHUNK_SIZE,
):
"""
计算图像-文本检索指标
指标:
- Image→Text Recall@1,5,10
- Text→Image Recall@1,5,10
- Mean Reciprocal Rank (MRR)
- Median Rank (MedR)
"""
# 计算相似度矩阵（分块以节省内存）
logits = logit_scale * (image_features @ text_features.T)
# 图像→文本检索
image2text_metrics = compute_recall(logits, topk=(1, 5, 10))
# 文本→图像检索
text2image_metrics = compute_recall(logits.T, topk=(1, 5, 10))
return {
"image2text_recall@1": image2text_metrics[0],
"image2text_recall@5": image2text_metrics[1],
"image2text_recall@10": image2text_metrics[2],
"text2image_recall@1": text2image_metrics[0],
"text2image_recall@5": text2image_metrics[1],
"text2image_recall@10": text2image_metrics[2],
}
```
### 完整评估函数 (train.py 第536-713行)
```python
def evaluate(task, data, epoch, args, tb_writer=None, tokenizer=None):
"""运行验证和零样本评估"""
metrics = {}
device = torch.device(args.device)
# 1. 零样本评估
zero_shot_metrics = zero_shot_eval_all(task, data, epoch, args, tokenizer)
metrics.update(zero_shot_metrics)
# 2. 验证集评估
if 'val' in data and (epoch % args.val_frequency == 0 or epoch == args.epochs):
task.eval()
dataloader = data['val'].dataloader
all_image_features = []
all_text_features = []
cumulative_loss = 0.0
with torch.inference_mode():
for batch in dataloader:
batch = task.prepare_batch(batch, device)
# 前向传播
model_out = task(batch)
# 收集特征
image_features = model_out["image_features"]
text_features = model_out["text_features"]
logit_scale = model_out["logit_scale"].mean()
all_image_features.append(image_features.cpu())
all_text_features.append(text_features.cpu())
# 计算损失
logits_per_image = logit_scale * (image_features @ text_features.T)
logits_per_text = logits_per_image.T
labels = torch.arange(batch_size, device=device)
loss = (F.cross_entropy(logits_per_image, labels) +
F.cross_entropy(logits_per_text, labels)) / 2
cumulative_loss += loss * batch_size
# 计算检索指标
val_metrics = get_clip_metrics(
image_features=all_image_features,
text_features=all_text_features,
logit_scale=logit_scale,
)
metrics.update(val_metrics)
# 3. 日志记录
_logger.info(f"Eval Epoch {epoch}: " + "\t".join([f"{k}: {v:.4f}" for k, v in metrics.items()]))
return metrics
```
### 图像检索示例
```python
def image_retrieval(query_image, database_images, model, preprocess, top_k=10):
"""
图像检索：从数据库中找到与查询最相似的图像
返回: (indices, similarities)
"""
# 1. 编码查询图像
query_feat = model.encode_image(preprocess(query_image).unsqueeze(0))
query_feat /= query_feat.norm(dim=-1, keepdim=True)
# 2. 编码数据库图像（批量）
db_feats = []
for img in database_images:
feat = model.encode_image(preprocess(img).unsqueeze(0))
feat /= feat.norm(dim=-1, keepdim=True)
db_feats.append(feat)
db_feats = torch.cat(db_feats, dim=0)
# 3. 计算相似度
similarities = (query_feat @ db_feats.T).squeeze(0)
# 4. Top-K
top_indices = similarities.topk(top_k).indices
top_similarities = similarities[top_indices]
return top_indices, top_similarities
```
### 文本-图像检索示例
```python
def text_image_search(query_text, image_dataset, model, tokenizer, preprocess, top_k=5):
"""使用文本查询搜索图像"""
# 1. 编码文本
text_tokens = tokenizer([query_text])
text_feat = model.encode_text(text_tokens)
text_feat /= text_feat.norm(dim=-1, keepdim=True)
# 2. 编码所有图像
image_feats = []
for img in image_dataset:
img_feat = model.encode_image(preprocess(img).unsqueeze(0))
img_feat /= img_feat.norm(dim=-1, keepdim=True)
image_feats.append(img_feat)
image_feats = torch.cat(image_feats, dim=0)
# 3. 计算相似度并排序
similarities = (text_feat @ image_feats.T).squeeze(0)
top_k_indices = similarities.topk(top_k).indices
return [(image_dataset[i], similarities[i].item()) for i in top_k_indices]
```
### 批量特征提取
```python
def extract_features_for_dataset(images, model, preprocess, batch_size=256, device='cuda'):
"""为整个数据集提取特征"""
model.eval()
model.to(device)
all_features = []
dataloader = DataLoader(images, batch_size=batch_size)
with torch.inference_mode():
for batch in dataloader:
batch = batch.to(device)
# 预处理
batch = torch.stack([preprocess(img) for img in batch])
# 编码
features = model.encode_image(batch)
features /= features.norm(dim=-1, keepdim=True)
all_features.append(features.cpu())
return torch.cat(all_features, dim=0)
```
### OpenCLIP vs OpenAI CLIP API对比
```python
# ===== OpenAI CLIP =====
import clip
model, preprocess = clip.load("ViT-B/32")
image_input = preprocess(Image.open("image.jpg")).unsqueeze(0)
text_input = clip.tokenize(["a dog", "a cat"])
image_features = model.encode_image(image_input)
text_features = model.encode_text(text_input)
logits_per_image, logits_per_text = model(image_input, text_input)
# ===== OpenCLIP =====
import open_clip
model, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
tokenizer = open_clip.get_tokenizer('ViT-B-32')
image_input = preprocess(Image.open("image.jpg")).unsqueeze(0)
text_input = tokenizer(["a dog", "a cat"])
image_features = model.encode_image(image_input)
text_features = model.encode_text(text_input)
logits_per_image, logits_per_text = model.get_logits(image_input, text_input)
# 或使用字典输出
model.output_dict = True
output = model(image=image_input, text=text_input)
# output = {"image_features": ..., "text_features": ..., "logit_scale": ...}
```
### 预训练模型列表
```python
>>> import open_clip
>>> open_clip.list_pretrained()[:10] # 显示前10个
[
('ViT-B-32', 'laion2b_s34b_b79k'),
('ViT-B-32', 'laion2b_s34b_b88k'),
('ViT-B-32', 'laion1b_s13b_b51k'),
('ViT-B-16', 'laion2b_s34b_b88k'),
('ViT-L-14', 'laion2b_s32b_b82k'),
('ViT-L-14', 'datacomp_xl_s13b_b90k'),
('ViT-H-14', 'laion2b_s32b_b79k'),
('convnext_base', 'laion2b_s34b_b88k'),
('convnext_large', 'laion2b_s29b_b131k'),
]
```
|
|
## 相关论文
### CLIP原论文
**Learning Transferable Visual Models From Natural Language Supervision**
- 作者：Alec Radford et al. (OpenAI)
- 年份：2021
- 亮点：大规模对比学习、零样本能力
### PMC-CLIP
**PMC-CLIP: Contrastive Learning of Medical Visual-Language Representation on PubMed Central**
- 针对医学图像-文本理解
- 在PMC开放数据上训练
- 使用BiomedBERT作为文本编码器
### 医学AI相关
- Medical image-text matching
- Clinical report generation
- Multi-modal medical AI
|
|
## 代码资源
### Hugging Face模型
- **PMC-CLIP权重**：https://huggingface.co/datasets/axiong/pmc_oa
- **BiomedBERT**：https://huggingface.co/microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract
### 示例代码
```python
# 完整推理示例
from transformers import AutoTokenizer, AutoModel
import torch
from PIL import Image
# 加载模型
tokenizer = AutoTokenizer.from_pretrained('microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract')
text_encoder = AutoModel.from_pretrained('microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract')
# 加载图像
image = Image.open('medical_image.jpg')
# 预处理
text_input = tokenizer(['chest X-ray'], return_tensors='pt')
image_input = preprocess(image).unsqueeze(0)
# 提取特征
with torch.no_grad():
image_feat = image_encoder(image_input)
text_feat = text_encoder(**text_input).pooler_output
# 计算相似度
similarity = (image_feat @ text_feat.T).item()
```
|