基于up主【【多模态】复现OpenAI的CLIP模型】 https://www.bilibili.com/video/BV13K421v7Ar/?share_source=copy_web&vd_source=62e4f3fd48d4ea21c7731322b8e5cb53 的代码 进行了修改 作出如下的学习笔记
一、代码展示与修改
1.CLIP 模型核心原理回顾
CLIP（Contrastive 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_train=True 表示默认加载训练集，设为 False 则加载测试集 def __init__(self, is_train=True): # 6. 调用父类的初始化方法（标准写法） super().__init__() # 7. 加载真正的 MNIST 数据 # root='./mnist/': 数据下载保存的路径 # train=is_train: 根据参数决定是下载训练集还是测试集 # download=True: 如果本地没有数据，自动从网上下载 self.ds = torchvision.datasets.MNIST('./mnist/', train=is_train, download=True) # 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_size=3, padding=1: 保持空间尺寸不变（不考虑 stride） self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, padding=1, stride=stride) # 7. 第一个批归一化层：加速收敛，稳定训练 self.bn1 = nn.BatchNorm2d(out_channels) # 8. 第二个卷积层：进一步提取特征，stride 固定为 1，不改变尺寸 self.conv2 = nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, padding=1, stride=1) # 9. 第二个批归一化层 self.bn2 = nn.BatchNorm2d(out_channels) # 10. 捷径连接（Shortcut/Projection）卷积层 # 当输入输出维度不一致时（通道数改变或尺寸改变），用 1x1 卷积进行匹配 self.conv3 = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, padding=0, stride=stride) # 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_channels=1, out_channels=16, stride=2) # 17. 第二个残差块：输入 16 通道，输出 4 通道，步长 2 # 输入 (14x14) -> 输出 (7x7) self.res_block2 = ResidualBlock(in_channels=16, out_channels=4, stride=2) # 18. 第三个残差块：输入 4 通道，输出 1 通道，步长 2 # 输入 (7x7) -> 输出 (4x4) self.res_block3 = ResidualBlock(in_channels=4, out_channels=1, stride=2) # 19. 全连接层（线性层）：将展平后的特征映射到最终的嵌入维度 # 此时特征图大小为 4x4x1 = 16，将其映射为 8 维向量 self.wi = nn.Linear(in_features=16, out_features=8) # 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. 创建一个模拟的输入张量：(Batch=1, Channel=1, Height=28, Width=28) # 对应一张 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_embeddings=10: 词表大小为 10，对应 MNIST 的 10 个数字类别（0-9） # embedding_dim=16: 将每个数字映射为一个 16 维的稠密向量 self.emb = nn.Embedding(num_embeddings=10, embedding_dim=16) # 7. 第一个全连接层：将 16 维向量映射到 64 维，用于提取更丰富的文本特征 self.dense1 = nn.Linear(in_features=16, out_features=64) # 8. 第二个全连接层：将 64 维向量降维回 16 维 self.dense2 = nn.Linear(in_features=64, out_features=16) # 9. 投影层（Projection）：将文本特征从 16 维映射到最终的 8 维空间 # 这里的 8 维与 ImgEncoder 输出的 8 维完全一致，这是 CLIP 跨模态对齐的关键 self.wt = nn.Linear(in_features=16, out_features=8) # 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(), lr=1e-3) ''' 训练模型主循环 ''' # 12. 定义总迭代次数：训练 10 万个 batch ITER_BATCH_COUNT = 100000 # 13. 定义批次大小：每次从数据集中抽取 64 个样本 BATCH_SIZE = 64 # 14. 定义目标类别数：MNIST 共有 10 种数字（0-9） TARGET_COUNT = 10 # 15. 创建数据加载器，开启 shuffle 打乱数据，num_workers=0 表示在主进程中加载数据 dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0) # 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, dim=0).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 列的子图，当前是第 i+1 个 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.svg：ResNet 残差块的示意图（若图像编码器用了 ResNet 结构，用于解释模块设计）。
README.md：项目说明（环境配置、运行步骤、功能介绍）。
3.代码修改
修改后的代码在资源绑定中，修改train.py中的两个地方
dataloader=DataLoader(dataset,batch_size=BATCH_SIZE,shuffle=True,num_workers=0) # 数据加载器在这里，将 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.代码修改后CTRL+s保存，否则依旧报错
