多模态-复现OpenAI的CLIP模型

基于up主【【多模态】复现OpenAI的CLIP模型】 https://www.bilibili.com/video/BV13K421v7Ar/?share_sourcecopy_webvd_source62e4f3fd48d4ea21c7731322b8e5cb53 的代码 进行了修改 作出如下的学习笔记一、代码展示与修改1.CLIP 模型核心原理回顾CLIPContrastive Language-Image Pre-training是 OpenAI 提出的跨模态预训练模型核心是通过「对比学习」让图像编码器和文本编码器的输出在特征空间中对齐——即“相似内容的图像和文本其特征向量距离更近不相似的则更远”。2.项目文件作用解析按功能分类1 数据层dataset.py作用定义 MNIST 数据集的加载、预处理逻辑。# 1. 导入 PyTorch 的数据集基类自定义数据集必须继承它 from torch.utils.data import Dataset # 2. 导入 torchvision 中的图像转换工具 # PILToTensor: 将 PIL 图片转换为 Tensor (张量) # Compose: 用于将多个转换步骤串联起来 from torchvision.transforms.v2 import PILToTensor, Compose # 3. 导入 torchvision 库主要用于获取内置的 MNIST 数据集 import torchvision # 4. 定义一个名为 MNIST 的类继承自 Dataset # 这是 PyTorch 加载数据的标准方式 class MNIST(Dataset): # 5. 初始化方法 # is_trainTrue 表示默认加载训练集设为 False 则加载测试集 def __init__(self, is_trainTrue): # 6. 调用父类的初始化方法标准写法 super().__init__() # 7. 加载真正的 MNIST 数据 # root./mnist/: 数据下载保存的路径 # trainis_train: 根据参数决定是下载训练集还是测试集 # downloadTrue: 如果本地没有数据自动从网上下载 self.ds torchvision.datasets.MNIST(./mnist/, trainis_train, downloadTrue) # 8. 定义图像预处理流程 # 这里只包含一步将图片转为 Tensor self.img_convert Compose([ PILToTensor(), ]) # 9. 返回数据集的总长度即图片的总数量MNIST 训练集通常是 60000 张 def __len__(self): return len(self.ds) # 10. 核心方法根据索引 index 获取某一张图片和对应的标签 def __getitem__(self, index): # 11. 从原始数据集中取出第 index 个样本 # img 是 PIL 格式的图片对象label 是整数0-9 img, label self.ds[index] # 12. 处理并返回数据 # self.img_convert(img): 将 PIL 图片转为 Tensor (数值范围 0-255) # / 255.0: 进行归一化将数值范围压缩到 0.0 - 1.0 之间利于模型训练 # 返回处理后的图片和原始标签 return self.img_convert(img) / 255.0, label # 13. 主程序入口只有直接运行此脚本时才会执行以下代码 if __name__ __main__: # 14. 导入绘图库 matplotlib import matplotlib.pyplot as plt # 15. 实例化我们定义的 MNIST 数据集类 ds MNIST() # 16. 获取数据集中的第 0 张图片及其标签 img, label ds[0] # 17. 在控制台打印这张图片的标签例如打印出 5 print(label) # 18. 使用 matplotlib 显示图片 # img.permute(1,2,0): 调整维度顺序。 # Tensor 默认是 (C, H, W) 即 (通道, 高, 宽)而 plt.imshow 需要 (H, W, C) plt.imshow(img.permute(1, 2, 0)) # 19. 弹出窗口展示图片 plt.show()2 模型层图像编码器 文本编码器CLIP 由双编码器组成项目中分别用独立文件实现img_encoder.py图像编码器作用将输入图像映射为固定维度的特征向量嵌入。# 1. 导入 PyTorch 的神经网络模块包含构建模型所需的各种层如卷积、线性层等 from torch import nn # 2. 导入 PyTorch 主库主要用于张量操作和随机数生成用于测试 import torch # 3. 导入 PyTorch 的函数式接口主要使用其中的激活函数如 ReLU import torch.nn.functional as F # 4. 定义残差块类继承自 nn.Module # 这是 ResNet 的核心组件用于解决深层网络梯度消失问题 class ResidualBlock(nn.Module): # 5. 初始化残差块的参数 def __init__(self, in_channels, out_channels, stride): super().__init__() # 6. 第一个卷积层负责主要的特征提取和下采样如果 stride 1 # kernel_size3, padding1: 保持空间尺寸不变不考虑 stride self.conv1 nn.Conv2d(in_channelsin_channels, out_channelsout_channels, kernel_size3, padding1, stridestride) # 7. 第一个批归一化层加速收敛稳定训练 self.bn1 nn.BatchNorm2d(out_channels) # 8. 第二个卷积层进一步提取特征stride 固定为 1不改变尺寸 self.conv2 nn.Conv2d(in_channelsout_channels, out_channelsout_channels, kernel_size3, padding1, stride1) # 9. 第二个批归一化层 self.bn2 nn.BatchNorm2d(out_channels) # 10. 捷径连接Shortcut/Projection卷积层 # 当输入输出维度不一致时通道数改变或尺寸改变用 1x1 卷积进行匹配 self.conv3 nn.Conv2d(in_channelsin_channels, out_channelsout_channels, kernel_size1, padding0, stridestride) # 11. 定义前向传播逻辑 def forward(self, x): # 12. 主路径卷积 - 归一化 - 激活 - 卷积 - 归一化 # 注意这里先做 BN 再做 ReLU是常见的变体写法 y F.relu(self.bn1(self.conv1(x))) y self.bn2(self.conv2(y)) # 13. 捷径路径直接通过 1x1 卷积调整输入 x 的维度使其能与 y 相加 z self.conv3(x) # 14. 残差连接将主路径输出 y 与捷径输出 z 相加最后再过一次激活函数 return F.relu(y z) # 15. 定义图像编码器类用于将 MNIST 图像转换为特征向量 class ImgEncoder(nn.Module): def __init__(self): super().__init__() # 16. 第一个残差块输入 1 通道灰度图输出 16 通道步长 2尺寸减半 # 输入 (28x28) - 输出 (14x14) self.res_block1 ResidualBlock(in_channels1, out_channels16, stride2) # 17. 第二个残差块输入 16 通道输出 4 通道步长 2 # 输入 (14x14) - 输出 (7x7) self.res_block2 ResidualBlock(in_channels16, out_channels4, stride2) # 18. 第三个残差块输入 4 通道输出 1 通道步长 2 # 输入 (7x7) - 输出 (4x4) self.res_block3 ResidualBlock(in_channels4, out_channels1, stride2) # 19. 全连接层线性层将展平后的特征映射到最终的嵌入维度 # 此时特征图大小为 4x4x1 16将其映射为 8 维向量 self.wi nn.Linear(in_features16, out_features8) # 20. 层归一化对最终的 8 维向量进行归一化常用于 CLIP 模型以计算余弦相似度 self.ln nn.LayerNorm(8) # 21. 定义前向传播 def forward(self, x): # 22. 依次通过三个残差块进行特征提取和下采样 x self.res_block1(x) x self.res_block2(x) x self.res_block3(x) # 23. 展平操作将 (Batch, 1, 4, 4) 变为 (Batch, 16) # 然后送入全连接层得到 8 维特征 x self.wi(x.view(x.size(0), -1)) # 24. 层归一化后返回最终结果 x self.ln(x) return x # 25. 主程序入口用于简单测试模型是否能跑通 if __name__ __main__: # 26. 实例化图像编码器 img_encoder ImgEncoder() # 27. 创建一个模拟的输入张量(Batch1, Channel1, Height28, Width28) # 对应一张 28x28 的灰度图 out img_encoder(torch.randn(1, 1, 28, 28)) # 28. 打印输出形状预期应为 torch.Size([1, 8]) print(out.shape)text_encoder.py文本编码器作用将文本标签如“数字3”映射为与图像编码器同维度的特征向量。# 1. 导入 PyTorch 的神经网络模块用于构建模型的各种层 from torch import nn # 2. 导入 PyTorch 主库用于张量操作如创建测试用的 tensor import torch # 3. 导入 PyTorch 的函数式接口这里主要用到 ReLU 激活函数 import torch.nn.functional as F # 4. 定义文本编码器类继承自 nn.Module # 作用是将 MNIST 数据集的文本标签0-9映射为与图像编码器维度相同的特征向量 class TextEncoder(nn.Module): # 5. 初始化方法 def __init__(self): super().__init__() # 6. 词嵌入层Embedding # num_embeddings10: 词表大小为 10对应 MNIST 的 10 个数字类别0-9 # embedding_dim16: 将每个数字映射为一个 16 维的稠密向量 self.emb nn.Embedding(num_embeddings10, embedding_dim16) # 7. 第一个全连接层将 16 维向量映射到 64 维用于提取更丰富的文本特征 self.dense1 nn.Linear(in_features16, out_features64) # 8. 第二个全连接层将 64 维向量降维回 16 维 self.dense2 nn.Linear(in_features64, out_features16) # 9. 投影层Projection将文本特征从 16 维映射到最终的 8 维空间 # 这里的 8 维与 ImgEncoder 输出的 8 维完全一致这是 CLIP 跨模态对齐的关键 self.wt nn.Linear(in_features16, out_features8) # 10. 层归一化LayerNorm对最终的 8 维向量进行归一化 # 确保文本特征和图像特征在同一个特征空间内便于后续计算余弦相似度 self.ln nn.LayerNorm(8) # 11. 定义前向传播逻辑 def forward(self, x): # 12. 输入整数标签如 3通过 Embedding 层转换为 16 维向量 x self.emb(x) # 13. 通过第一个全连接层并使用 ReLU 激活函数引入非线性 x F.relu(self.dense1(x)) # 14. 通过第二个全连接层同样使用 ReLU 激活 x F.relu(self.dense2(x)) # 15. 通过投影层将特征维度从 16 压缩到 8 x self.wt(x) # 16. 进行层归一化返回最终的 8 维文本特征向量 x self.ln(x) return x # 17. 主程序入口用于测试文本编码器是否能正常运行 if __name__ __main__: # 18. 实例化文本编码器 text_encoder TextEncoder() # 19. 创建一个模拟的输入张量包含 0-9 的 10 个数字标签 x torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 0]) # 20. 将 10 个标签送入文本编码器得到 10 个对应的 8 维特征向量 y text_encoder(x) # 21. 打印输出形状预期结果应为 torch.Size([10, 8]) print(y.shape)3训练层train.py作用定义 CLIP 的训练流程核心是对比损失Contrastive Loss。# 1. 导入 PyTorch 核心库 import torch # 2. 导入之前自定义的 MNIST 数据集类 from dataset import MNIST # 3. 导入之前定义的 CLIP 模型类包含图像编码器和文本编码器 from clip import CLIP # 4. 导入 PyTorch 的函数式接口这里主要用到交叉熵损失函数 import torch.nn.functional as F # 5. 导入数据加载器用于批量加载数据 from torch.utils.data import DataLoader # 6. 导入操作系统模块用于文件的重命名安全保存模型 import os # 7. 设置运行设备如果有 NVIDIA GPU 则使用 cuda否则使用 cpu DEVICE cuda if torch.cuda.is_available() else cpu # 8. 实例化 MNIST 数据集 dataset MNIST() # 9. 实例化 CLIP 模型并将其移动到指定设备GPU/CPU上 model CLIP().to(DEVICE) # 10. 尝试加载之前保存的模型权重断点续训 try: # 从 model.pth 文件中读取状态字典并加载到模型中 model.load_state_dict(torch.load(model.pth)) except: # 如果文件不存在或加载失败则忽略错误从头开始训练 pass # 11. 定义 Adam 优化器用于更新模型参数学习率设为 0.001 optimzer torch.optim.Adam(model.parameters(), lr1e-3) 训练模型主循环 # 12. 定义总迭代次数训练 10 万个 batch ITER_BATCH_COUNT 100000 # 13. 定义批次大小每次从数据集中抽取 64 个样本 BATCH_SIZE 64 # 14. 定义目标类别数MNIST 共有 10 种数字0-9 TARGET_COUNT 10 # 15. 创建数据加载器开启 shuffle 打乱数据num_workers0 表示在主进程中加载数据 dataloader DataLoader(dataset, batch_sizeBATCH_SIZE, shuffleTrue, num_workers0) # 16. 开始迭代训练循环 for i in range(ITER_BATCH_COUNT): # 17. 内部死循环不断抽取 batch直到抽到一个包含全部 10 种数字的 batch 为止 while True: # 从数据加载器中获取下一个 batch 的图片和标签 imgs, labels next(iter(dataloader)) # 统计当前 batch 中不重复的标签数量如果小于 10说明没凑齐所有数字 if torch.unique(labels).shape[0] TARGET_COUNT: continue # 继续抽取下一个 batch # 凑齐了 10 种数字后开始从中挑选出刚好 10 个不重复的样本 target set() # 用集合记录已经选中的数字 indexes [] # 记录选中样本的索引 # 遍历当前 batch 的 64 个样本 for j in range(BATCH_SIZE): # 如果当前数字已经被选中过则跳过 if labels[j].item() in target: continue # 将新数字加入集合记录其索引 target.add(labels[j].item()) indexes.append(j) # 如果已经选够了 10 个不同的数字则跳出循环 if len(target) TARGET_COUNT: break # 根据选出的索引提取出包含 10 个不同数字的图片和标签 imgs imgs[indexes] labels labels[indexes] break # 成功凑齐 10 个样本跳出 while True 循环 # 18. 将挑选好的 10 个图片和标签送入 CLIP 模型得到相似度矩阵 (logits) logits model(imgs.to(DEVICE), labels.to(DEVICE)) # 19. 构造真实标签Ground Truth对角线 [0, 1, 2, ..., 9] # 表示第 0 张图对应第 0 个文本第 1 张图对应第 1 个文本以此类推 targets torch.arange(0, TARGET_COUNT).to(DEVICE) # 20. 计算图像到文本方向的交叉熵损失 loss_i F.cross_entropy(logits, targets) # 21. 计算文本到图像方向的交叉熵损失 # permute(1,0) 将矩阵转置实现双向对齐 loss_t F.cross_entropy(logits.permute(1, 0), targets) # 22. 计算最终的对比损失取双向损失的平均值 loss (loss_i loss_t) / 2 # 23. 清空上一步累积的梯度PyTorch 默认会累加梯度 optimzer.zero_grad() # 24. 反向传播计算当前 loss 对模型参数的梯度 loss.backward() # 25. 优化器根据梯度更新模型参数 optimzer.step() # 26. 每训练 1000 次迭代打印一次当前的迭代次数和损失值 if i % 1000 0: print(iter:{},loss:{}.format(i, loss)) # 27. 安全保存模型权重 # 先保存到临时文件 .model.pth torch.save(model.state_dict(), .model.pth) # 28. 将临时文件替换为正式的 model.pth # 这样做是为了防止在保存过程中程序崩溃导致原模型文件损坏 os.replace(.model.pth, model.pth)4推理层inference.py作用利用训练好的模型做跨模态检索/分类。# 1. 导入之前自定义的 MNIST 数据集类 from dataset import MNIST # 2. 导入绘图库用于可视化展示图片和相似度结果 import matplotlib.pyplot as plt # 3. 导入 PyTorch 核心库 import torch # 4. 导入之前定义的 CLIP 模型类 from clip import CLIP # 5. 导入 PyTorch 的函数式接口本脚本中虽未直接使用但属于常规导入 import torch.nn.functional as F # 6. 设置运行设备优先使用 GPU否则使用 CPU DEVICE cuda if torch.cuda.is_available() else cpu # 7. 实例化 MNIST 数据集 dataset MNIST() # 8. 实例化 CLIP 模型并移动到指定设备 model CLIP().to(DEVICE) # 9. 加载训练好的模型权重文件 model.load_state_dict(torch.load(model.pth)) # 10. 将模型切换到评估模式eval # 这会关闭 Dropout 和 BatchNorm 的训练行为确保推理结果稳定 model.eval() 1、对图片分类Zero-Shot Classification # 11. 从数据集中取出第 0 张图片及其真实标签 image, label dataset[0] # 12. 打印这张图片的真实分类例如 5 print(正确分类:, label) # 13. 显示这张原始图片 # permute(1,2,0) 将 Tensor 维度从 (C,H,W) 转换为 (H,W,C) 以适应 matplotlib plt.imshow(image.permute(1, 2, 0)) plt.show() # 14. 构造包含 0-9 所有类别的文本标签张量 targets torch.arange(0, 10) # 15. 将 1 张图片和 10 个文本标签一起送入 CLIP 模型 # unsqueeze(0) 在 batch 维度增加一维变成 (1, 1, 28, 28) # 模型会输出一个 (1, 10) 的相似度矩阵logits logits model(image.unsqueeze(0).to(DEVICE), targets.to(DEVICE)) # 16. 打印模型输出的原始相似度分数 print(logits) # 17. 找出相似度分数最高的类别索引作为 CLIP 的预测结果 # argmax(-1) 在最后一个维度10个类别上找最大值 print(CLIP分类:, logits.argmax(-1).item()) 2、图像相似度检索Image-to-Image Retrieval # 18. 准备用于检索的候选图片列表和对应的标签列表 other_images [] other_labels [] # 19. 从数据集中遍历第 1 到第 100 张图片作为候选图库 for i in range(1, 101): other_image, other_label dataset[i] other_images.append(other_image) other_labels.append(other_label) # 20. 将 100 张候选图片堆叠成一个张量送入图像编码器提取特征 # torch.stack 将列表变为 (100, 1, 28, 28) 的张量 # 输出 other_img_embs 的形状为 (100, 8) other_img_embs model.img_enc(torch.stack(other_images, dim0).to(DEVICE)) # 21. 将当前测试图片也送入图像编码器提取特征 # 输出 img_emb 的形状为 (1, 8) img_emb model.img_enc(image.unsqueeze(0).to(DEVICE)) # 22. 通过矩阵乘法计算当前图片与 100 张候选图片的余弦相似度点积 # (1, 8) (8, 100) (1, 100) # 注代码中变量名拼写为 logtis应为 logits logtis img_emb other_img_embs.T # 23. 找出相似度最高的前 5 个候选图片 # values 是前 5 高的相似度分数indexs 是对应的图片索引 values, indexs logtis[0].topk(5) # 24. 创建一个 15x15 英寸的大画布用于展示检索结果 plt.figure(figsize(15, 15)) # 25. 遍历最相似的 5 张图片并依次绘制 for i, img_idx in enumerate(indexs): # 26. 创建 1 行 5 列的子图当前是第 i1 个 plt.subplot(1, 5, i 1) # 27. 显示这张最相似的图片 plt.imshow(other_images[img_idx].permute(1, 2, 0)) # 28. 在图片上方添加标题显示其真实标签 plt.title(other_labels[img_idx]) # 29. 关闭坐标轴显示使图片更美观 plt.axis(off) # 30. 弹出窗口展示最终的 5 张最相似图片 plt.show()5辅助文件clip.py CLIP 模型的整体封装整合图像/文本编码器定义前向传播逻辑。refs存放参考文献、技术文档如 CLIP 原论文链接。clip-model.png / similar_image.png / image.png可视化素材模型架构图、相似图像示例、测试用图像。resnet-block.svgResNet 残差块的示意图若图像编码器用了 ResNet 结构用于解释模块设计。README.md项目说明环境配置、运行步骤、功能介绍。3.代码修改修改后的代码在资源绑定中修改train.py中的两个地方dataloaderDataLoader(dataset,batch_sizeBATCH_SIZE,shuffleTrue,num_workers0) # 数据加载器在这里将 num_workers 改为 0并将 persistent_workers 改为 False (或者直接删掉 persistent_workers 这个参数)原因多进程冲突DataLoader 默认会启动多个子进程num_workers 0来并行加载数据。将 num_workers 设置为 0最快速有效这是最简单直接的解决方法。它会禁用多进程让数据加载在主进程中完成。虽然速度会慢一些但能立即排除多进程相关的所有问题。既然你选择了不使用 worker数量为 0那就没有进程可以“保持存活”了。PyTorch 检测到这个逻辑冲突所以抛出了 ValueError。二、运行结果1.运行train.py后 终端显示程序正在正常运行Loss 在不断下降说明模型正在学习2.运行inference.py后显示识别正确终端底部显示正确分类: 5说明系统准确判断出这是一张数字5。三、注意事项1.检查torchvision版本太旧不能运行2.代码修改后CTRLs保存否则依旧报错