从零开始搭建CLIP模型实现基于文本的图像检索

CLIP原理简介

   论文链接，源码链接
  CLIP模型由OpenAI在2021年提出，利用双Decoder（Dual Encoder）的架构来学习图像和文本之间的对应关系，是多模态大模型的开创之作，为后续许多高效的多模态模型的提出打下基础。CLIP是一个预练习模型(Pre-trained Model)，在学习到图像–文本特性之间的关联后可以迁移到各种下游任务中，如图像分类，文本引导图像分割和目标检测，图像文本检索等。由于模型学习到的是文本语义和图像语义之间的关联，使得其zero-shot能力非常强盛，根据论文中的描述，CLIP在许多数据集上zero-shot的结果乃至超越了许多练习好的模型的效果。CLIP的练习范式如下：

CLIP的布局非常简单，数据集包含大量的图像文本对，图像经过图像编码器得到图像特性，文本经过文本编码器得到文本特性，将图像特性和文本特性按照数据集中的对应关系进行配对，不配对的特性给予惩罚，从上图中可以看出，我们希望矩阵中蓝色的值趋近于1，别的值趋近于0，采用对比学习的方式对模型进行练习，算法的伪代码如下：

从损失函数中可以看出，分别对特性对比矩阵的行和列进行交叉熵损失函数计算，并取平均得到终极的loss。图像编码器一样寻常有两种选择：ResNet和ViT；文本编码器采用Transformer Encoder，均是各自范畴中优秀的特性提取网络。
CLIP的推理范式如下：

在推理阶段，图像编码器中输入图像获取图像特性，文本编码器中输入文本获取文本特性，将图像特性向量和文本特性向量的转置相乘得到每张图像对每个文本的特性相似度，相似度最高的文本即描述了该图像中物体所属的种别。
代码实现

   Flickr8k数据集下载，提取码：fbfz
DistilBert模型文件下载
  我的运行环境：
 CUDA 11.8
 pytorch 2.2.2
 transformers 4.44.0 # 用于从HuggingFace上加载预练习模型

---

数据集预览：

   图片示例  

   文本示例  

---

由于作者的显卡算力有限，选取Flickr8k数据集进行模型练习，其中包含8k个图像文本对，其中一张图像对应5条文本。图像编码器采用ResNet50，直接从timm库中导入；文本编码器采用DistilBert，即轻量化的Bert模型，从HuggingFace上下载。闲话少说，小二，上菜！

1. ### 模型参数配置 ###
2. import argparse
3. from dataclasses import dataclass
5. parser = argparse.ArgumentParser(description="CLIP from zero")
6. parser.add_argument("--image_dir", default="user/Flickr8k/Images", help='path to image folder')  # 存放图像的文件路径
7. parser.add_argument("--caption_dir", default="user/Flickr8k", help='path to caption folder')  # 存放文本的文件路径
8. parser.add_argument("--weight_dir", default='user/checkpoints', help='path to save output weight')  # 存放训练权重的文件路径
9. args = parser.parse_args()
11. @dataclass
12. class CLIPConfig:
13.     image_path: str = args.image_dir  # 图像存放路径
14.     image_size: int = 224  # resize后的图像尺寸，便于构建Dataloader
15.     caption_path: str = args.caption_dir  # 文本存放路径
16.     batch_size: int = 8  # 一个批次中的数据数量
17.     epochs: int = 3  # 训练世代
19.     image_encoder_model: str = "resnet50"  # 图像编码器的名称
20.     image_embedding_dim: int = 2048  # 图像特征的维度
21.     text_encoder_model: str = "distilbert-base-uncased"  # 文本编码器的名称
22.     text_embedding_dim: int = 768  # 文本特征的维度
23.     text_tokenizer: str = text_encoder_model  # 文本分词器模型的名称
24.     max_length: int = 200  # 文本编码器可输入的最长文本长度
26.     pretrained: bool = False  # 是否加载预训练好的编码器
27.     trainable: bool = True  # 在训练过程中是否更新编码器的参数
28.     temperature: float = 1.0  # 计算loss时的正则化系数
30.     proj_dim: int = 256  # 图像特征和文本特征统一后的维度
31.     dropout_rate: float = 0.1  # dropout系数，避免过拟合
34. ### 载入数据集并初始化 ###
35. import torch
36. from torch.utils.data import Dataset, DataLoader
37. from transformers import DistilBertModel, DistilBertConfig, DistilBertTokenizer
38. import albumentations as A
39. import pandas as pd
40. import cv2
42. class CLIPDataset(Dataset):
43.     def __init__(self, config, image_path, caption_path, transforms=True):
44.         """
45.         图片文件名和标题的长度必须相同
46.         如果一个图片对应多个标题，该图片文件名需要重复多次
47.         """
48.         self.image_path = image_path  # 图像路径
49.         self.caption_path = caption_path  # 文本路径
50.         self.dataframe = pd.read_csv(f"{self.caption_path}/captions.csv")  # 读取文本
51.         self.tokenizer = DistilBertTokenizer.from_pretrained(config.text_tokenizer)  # 载入分词器
53.         self.image_filenames = self.dataframe["image"].values  # 获取图像文件名
54.         self.captions = list(self.dataframe["caption"].values)   # 获取图像对应的描述文本
55.         self.encoded_captions = self.tokenizer(self.captions,
56.                                                padding=True,
57.                                                truncation=True,
58.                                                max_length=config.max_length)  # 文本分词
59.         self.transforms = transforms  # 对输入图像进行预处理
61.     def __getitem__(self, idx):  # 获取数据集中第idx个数据，其中包含图片名称和对应的标题(可能不止一个)
62.         item = {
63.             key: torch.tensor(values[idx]) for key, values in self.encoded_captions.items()
64.         }
66.         image = cv2.imread(f"{self.image_path}/{self.image_filenames[idx]}")  # 获取原始图像
67.         image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
68.         if self.transforms:
69.             image = self.get_transforms(mode="train")(image=image)["image"]  # 对图像进行预处理
71.         item["image"] = torch.tensor(image).permute(2, 0, 1).float()  # 将图片转换为tensor格式，并调整为RGB顺序
72.         item["caption"] = self.captions[idx]  # 获取标题
74.         return item
76.     def __len__(self):
77.         return len(self.captions)  # 获取文本长度
79.     def get_transforms(self, mode="train"):
80.         if mode == "train":
81.             return A.Compose(
82.                 [
83.                     A.Resize(config.image_size, config.image_size, always_apply=True),  # 对图像进行resize
84.                     A.Normalize(max_pixel_value=255.0, always_apply=True)  # 对像素值进行归一化
85.                 ]
86.             )
88. ### 图像编码器 ###
89. import torch.nn as nn
90. import timm
92. class ImageEncoder(nn.Module):
93.     """
94.     图像编码器，采用ResNet50
95.     """
96.     def __init__(self, config):
97.         super().__init__()
98.         self.model = timm.create_model(config.image_encoder_model,
99.                                        pretrained=config.pretrained,
100.                                        num_classes=0, global_pool="avg")  # 创建ResNet50
102.         for p in self.model.parameters():
103.             p.requires_grad = config.trainable  # 设置参数可训练
105.     def forward(self, x):
106.         image_encoded = self.model(x)  # 获得图像特征编码，形状为[batch_size, image_embedding_dim]
107.         return image_encoded
109. ### 文本编码器 ###
110. class TextEncoder(nn.Module):
111.     """
112.     文本编码器，采用DistilBERT
113.     """
114.     def __init__(self, config):
115.         super().__init__()
116.         if config.pretrained:
117.             self.model = DistilBertModel.from_pretrained(config.text_encoder_model)  # 导入下载好的模型文件
118.         else:
119.             self.model = DistilBertModel(DistilBertConfig())
121.         for p in self.model.parameters():
122.             p.requires_grad = config.trainable  # 设置参数可训练
124.         self.target_token_idx = 0
125.    
126.     # 提取出和图像对应的文本特征向量
127.     def forward(self, input_ids, attention_mask):
128.         output = self.model(input_ids=input_ids, attention_mask=attention_mask)
129.         text_encoded = output.last_hidden_state[:, self.target_token_idx, :]  # [batch_size, text_embedding_dim]
130.         return text_encoded
132. ### 投影层 (MLP) ###
133. class ProjectionHead(nn.Module):
134.     """
135.     将图像编码和文本编码映射到相同维度
136.     """
137.     def __init__(self, config, input_embedding_dim):
138.         super().__init__()
139.         self.proj = nn.Linear(input_embedding_dim, config.proj_dim)
140.         self.act_fn = nn.GELU()
141.         self.fc = nn.Linear(config.proj_dim, config.proj_dim)
142.         self.dropout = nn.Dropout(config.dropout_rate)
143.         self.layer_norm = nn.LayerNorm(config.proj_dim)
145.     def forward(self, x):
146.         x_proj = self.proj(x)
147.         x = self.act_fn(x_proj)
148.         x = self.fc(x)
149.         x = self.dropout(x)
150.         x = x + x_proj
151.         x = self.layer_norm(x)
153.         return x
155. ### 定义损失函数 ###
156. def cross_entropy(logits, labels, reduction='none'):
157.     log_softmax = nn.LogSoftmax(dim=-1)
158.     loss = (-labels * log_softmax(logits)).sum(dim=1)
159.     if reduction == 'mean':
160.         return loss.mean()
161.     else:
162.         return loss.sum()
164. ### 模型主体 ###
165. import torch.nn.functional as F
167. class CLIP(nn.Module):
168.     def __init__(self, config):
169.         super().__init__()
170.         self.image_encoder = ImageEncoder(config)  # 实例化图像编码器
171.         self.text_encoder = TextEncoder(config)  # 实例化文本编码器
172.         self.image_proj = ProjectionHead(config, config.image_embedding_dim)  # 图像特征投影
173.         self.text_proj = ProjectionHead(config, config.text_embedding_dim)  # 文本特征投影
174.         self.temperature = config.temperature
176.     def forward(self, batch):
177.         image_features = self.image_encoder(batch["image"])  # 图像编码
178.         
179.         # 文本编码，tokenizer处理后的文本序列自带input_ids和attention_mask
180.         text_features = self.text_encoder(batch["input_ids"], batch["attention_mask"])
182.         image_embeddings = self.image_proj(image_features)  # 图像特征投影
183.         text_embeddings = self.text_proj(text_features)  # 文本特征投影
185.         logits = (text_embeddings @ image_embeddings.T) / self.temperature  # tensor形状为[batch_size, batch_size]
186.         images_similarity = image_embeddings @ image_embeddings.T  # tensor形状为[batch_size, batch_size]
187.         text_similarity = text_embeddings @ text_embeddings.T  # tensor形状为[batch_size, batch_size]
189.         # 软标签，不配对的位置设置为较小的数，而非0
190.         labels = F.softmax((images_similarity + text_similarity) / 2 * self.temperature, dim=-1)  
191.         
192.         loss_T = cross_entropy(logits, labels)  # 计算文本损失
193.         loss_I = cross_entropy(logits.T, labels.T)  # 计算图像损失
194.         total_loss = (loss_T + loss_I) / 2  # 对比学习平均损失
196.         return total_loss, logits
198. ### 训练函数 ###
199. def train(model, optimizer, scheduler, train_loader, device):
200.     model.train()  # 模型设置为训练模式
201.     train_loss = 0
202.     train_loader = tqdm(train_loader, total=len(train_loader))  # 显示训练进度条
203.     cnt = 0
204.     for batch in train_loader:
205.         # print(batch.keys())
206.         cnt += 1
207.         batch = {k: v.to(device) for k, v in batch.items() if k != "caption"}  # 将dataloader中一个batch的数据转换为字典形式
208.         loss, _ = model(batch)
209.         optimizer.zero_grad()
210.         loss.backward()
211.         optimizer.step()
212.         scheduler.step(metrics=loss.item())  # 根据上次训练的损失更新学习率
214.         train_loss += loss.item()
216.         # 训练100个batch显示一次loss
217.         if cnt % 100 == 0:
218.             print(f' ==> Epoch: {epoch + 1}, Batch: {cnt}, Loss: {loss.item():.4f}')
220.     return train_loss / len(train_loader)  # 平均训练损失
222. ### 测试函数 ###
223. def eval(model, val_loader, device):
224.     model.eval()  # 模型设置为测试模式
225.     val_loss = 0
226.     val_loader = tqdm(val_loader, total=len(val_loader))
227.     with torch.no_grad():
228.         for batch in val_loader:
229.             batch = {k: v.to(device) for k, v in batch.items() if k != "caption"}
230.             loss, _ = model(batch)
231.             val_loss += loss.item()
233.     return val_loss / len(val_loader)  # 平均测试损失
235. if __name__ == '__main__':
236.     config = CLIPConfig()  # 实例化配置信息
237.     model = CLIP(config)  # 实例化CLIP模型
238.     device = "cuda" if torch.cuda.is_available() else "cpu"
239.     model = model.to(device)
241.     # 查看模型的总参数量
242.     total_params = sum(p.numel() for p in model.parameters())
243.     print(f"Total parameters: {total_params / 1e6} M")
245.     optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-3)
246.     scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", patience=2, factor=0.5)
248.     dataset = CLIPDataset(config, args.image_dir, args.caption_dir)  # 读取并预处理数据
249.     train_dataset, val_dataset = torch.utils.data.random_split(dataset, [0.8, 0.2])  # 80%为训练数据，20%为测试数据
250.     dataloader = DataLoader(dataset, batch_size=config.batch_size, shuffle=False)
251.     train_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True)
252.     val_loader = DataLoader(val_dataset, batch_size=config.batch_size, shuffle=False)
254.     # 开始训练
255.     best_loss = float("inf")
256.     for epoch in range(config.epochs):
257.         print(f"Epoch: {epoch + 1}")
258.         train_loss_avg = train(model, optimizer, scheduler, train_loader, device)
259.         val_loss_avg = eval(model, val_loader, device)
261.         if val_loss_avg < best_loss:
262.             best_loss = val_loss_avg
263.             torch.save(model.state_dict(), f'{args.weight_dir}' + f'/CLIP_{epoch}.pth')
264.             print("Best model saved!")
266.     # 图像文本检索推理并可视化
267.     # dataframe = pd.read_csv(f"{config.caption_path}/captions.csv")
268.     # tokenizer = DistilBertTokenizer.from_pretrained(config.text_tokenizer)
269.     # model.load_state_dict(torch.load(f'{args.weight_dir}' + f'/CLIP_1.pth', map_location=device))
270.     # model.eval()
271.     #
272.     # image_embeddings = []
273.     # with torch.no_grad():
274.     #     for batch in tqdm(dataloader):
275.     #         image_features = model.image_encoder(batch["image"].to(device))  # 获取图像特征
276.     #         cur_image_embeddings = model.image_proj(image_features)  # [batch_size, proj_dim]  # 图像特征投影
277.     #         image_embeddings.append(cur_image_embeddings)  # 将一个batch的图像特征保存
278.     #
279.     # image_embeddings = torch.cat(image_embeddings, dim=0)  # [image_number, proj_dim]
280.     # input_query = "two dogs sitting on the grass"  # 输入文本
281.     # image_filenames = dataframe["image"].values  # 待检索的图片
282.     #
283.     # encoded_query = tokenizer([input_query])  # 对输入文本进行分词
284.     # batch = {key: torch.tensor(values).to(device) for key, values in encoded_query.items()}
285.     #
286.     # with torch.no_grad():
287.     #     text_features = model.text_encoder(batch["input_ids"], batch["attention_mask"])  # 获取文本特征
288.     #     text_embeddings = model.text_proj(text_features)  # 文本特征投影，与图像特征维度一致
289.     #
290.     # image_embeddings_n = F.normalize(image_embeddings, dim=-1)  # [image_number, proj_dim]
291.     # text_embeddings_n = F.normalize(text_embeddings, dim=-1)  # [1, proj_dim]
292.     # dot_similarity = text_embeddings_n @ image_embeddings_n.T  # 输入文本的特征和数据集中每张图像特征之间的相似度
293.     #
294.     # values, indices = torch.topk(dot_similarity.squeeze(0), k=45)  # 获取前45个相似度最高的图像
295.     # matches = [image_filenames[idx] for idx in indices[::5]]  # 获取对应的图像文件名(9张图像)
296.     #
297.     # f, axes = plt.subplots(3, 3, figsize=(10, 10))
298.     # f.suptitle(f"Retrieving text: {input_query}")  # 设置主标题
299.     # for match, ax in zip(matches, axes.flatten()):  # 显示检索出的图像
300.     #     image = cv2.imread(f"{args.image_dir}/{match}")
301.     #     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
302.     #     ax.imshow(image)
303.     #     ax.axis("off")
304.     #
305.     # plt.show()

复制代码

理想结果：

参考链接

https://towardsdatascience.com/simple-implementation-of-openai-clip-model-a-tutorial-ace6ff01d9f2/

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。