学习周报二十
摘要
本周主要进行了扩散模型代码的深入复现与多模态技术学习。在代码实践方面,系统研究了基于Stable Diffusion的CLIP增强模型,重点分析了模型架构中的视觉特征投影、UNet噪声预测机制和零样本分类器实现;解决了Flash Attention安装中的技术难题并编写了详细教程;在理论学习方面,深入探索了DeepSeek-OCR的技术演进,包括Vary模型的架构设计和视觉压缩文本的创新方法。通过代码级解析与技术文档编写,建立了从模型实现到优化部署的完整认知。
Abstract
This week focused on in-depth reproduction of diffusion model code and multimodal technology learning. In code practice, systematically studied the CLIP enhancement model based on Stable Diffusion, with emphasis on visual feature projection, UNet noise prediction mechanism, and zero-shot classifier implementation in the model architecture. Resolved technical challenges in Flash Attention installation and compiled detailed tutorials. In theoretical learning, deeply explored the technical evolution of DeepSeek-OCR, including the architectural design of Vary model and innovative methods for visual text compression. Through code-level analysis and technical documentation, established a complete cognition from model implementation to optimization deployment.
1、代码复现
本周学习了一篇使用扩散模型提升Clip细腻度的论文,对其主要的model代码进行了详细学习,关键部分做了标注学习。
from dataclasses import dataclass
from typing import Optional
from dataclasses import dataclass
import torch
import torch.nn as nn
from einops import rearrange
from transformers.utils import ModelOutput
from transformers.modeling_utils import PreTrainedModel
from .build import load_sd_model, load_clip_model_OpenAICLIP
from .utils import initiate_time_steps, prepare_class_text_embeddings
import torchvision.transforms as transforms
import clip@dataclass
class SDOutput(ModelOutput):loss: Optional[torch.FloatTensor] = Noneclass SDModel(PreTrainedModel):def __init__(self,config = None,):super().__init__(config)# 加载Stable Diffusion组件# image_renormalizer:把图像像素变换到 VAE 期望的范围(比如 [0,1] → [-1,1])self.model_id, self.pipe, self.vae, self.tokenizer, self.text_encoder, self.unet, self.scheduler, self.image_renormalizer = load_sd_model(config)# .eval() 让这些模块处于推理态(关掉 Dropout 等);是否更新参数由你外部的优化器决定self.text_encoder.eval()self.vae.eval()# 扩散模型的核心网络,输入是带噪的 z_t 和时间步 t,在条件 context 下预测噪声self.unet.eval()self.pattern_dictionary={'None':['']}self.config.actual_bs = len(self.pattern_dictionary[self.config.visual_pattern])# 加载OpenAI CLIP模型self.class_model = load_clip_model_OpenAICLIP(config)self.class_model.eval()self.config = configdiscrimi_size = self.config.clip_image_sizeself.resize_transform_discrimi = transforms.Resize((discrimi_size, discrimi_size))# 关键:视觉特征投影层self.visual_proj = nn.Linear(768, 1024)# CLIP特征 → 扩散模型条件空间# 真正做“特征 + 分类”的函数# image: 图像张量# classes: 类别索引列表,如果是 None,就表示用所有类别def classify(self, image, classes):# CLIP提取特征+打分# logits 由 final_fc(image_features) 得到,形状 [B, C]:每张图对每个类的一个分数。# logits 形状是 [B, C],每一行是一张图在 C 个类别上的打分image_features, logits = self.class_model(image)# logits[:, classes] 是 PyTorch 的高级索引:取所有行(:),但只取第 classes 指定的那些列。结果形状会变成 [B, len(classes)]if classes is not None:logits = logits[:, classes]# 将打分变为概率probs = logits.softmax(-1)# 拿每张图概率最大的那个类的索引max_idx = probs.argmax(-1)# probs.shape[-1] 表示 probs 的最后一个维度的长度,在这里,probs 的形状是 [B, C],所以 probs.shape[-1] = C,即类别总数# 如果 adapt_topk == -1,就表示“取所有类别”,否则,就取配置文件中指定的前 K 个概率最大的类K = probs.shape[-1] if self.config.tta.adapt_topk == -1 else self.config.tta.adapt_topk'''取出前 K 个概率最高的类别的索引probs.argsort(descending=True)argsort() 会返回从大到小排序后,每个元素在原数组中的索引。descending=True 表示按概率从高到低排序。'''topk_idx = probs.argsort(descending=True)[:, :K]# 把“在子类别中的索引”转换回“在原始类别列表中的索引”if classes is not None:classes = torch.tensor(classes).to(logits.device)max_class_idx = classes[max_idx.flatten()].view(max_idx.shape)topk_class_idx = classes[topk_idx.flatten()].view(topk_idx.shape)else:max_class_idx, topk_class_idx = max_idx, topk_idx'''CLIP 的工作结果就是:一张图片 → 一个语义向量 image_features([B,768])一张图片 → 对每个类的打分 logits([B,C]),再变成概率 probs([B,C])每张图选出最可能的 K 个类 topk_idx([B,K])'''return image_features, logits, topk_idx, max_class_idx, topk_class_idx# 从一张干净图像的潜变量(x_start)开始;# 逐步往里面加噪声(得到 x_t);# 然后训练一个网络(UNet)去预测出“这一步加了多少噪声”(pred_noise);# 反复迭代,最终可以从纯噪声还原出图像。def _unet_pred_noise(self, x_start, t, noise, context):# x_start: 原始图像的潜空间向量(latent),来自 VAE 编码器# t: 时间步(例如 0~999),代表“扩散到第几步”,噪声强度随 t 增加而增加# noise: 真实添加的噪声# context: 条件信息(一般是文本特征,比如 CLIP 编码的 prompt 向量),用于“指导”生成_,c,h,w = x_start.shapedevice = t.devicent = t.shape[0]x_start = x_start.unsqueeze(1)x_start = x_start.expand(-1, nt//x_start.shape[0], -1, -1, -1)x_start = x_start.reshape(-1,c,h,w)# alphas_cumprod 是扩散调度器里保存的累计噪声衰减系数,表示每个时间步 t 的“保留原图成分”的比例。# 例如:t=0 → α_cumprod ≈ 1.0 (几乎全是原图) t=999 → α_cumprod ≈ 0.0 (几乎全是噪声)alphas_cumprod = self.scheduler.alphas_cumprod.to(device)# 人工添加噪声,公式是:x_t = √(α_cumprod) * x_0 + √(1 - α_cumprod) * noise# noised_latent即x_t x_start即x_0noised_latent = (x_start * (alphas_cumprod[t]**0.5).view(-1, 1, 1, 1).to(device)+ noise * ((1 - alphas_cumprod[t])**0.5).view(-1, 1, 1, 1).to(device))# 调用 UNet 预测噪声pred_noise = self.unet(noised_latent, t, encoder_hidden_states=context.expand(nt, -1, -1)).samplereturn pred_noise# 根据类别名(classnames)和模板(templates),生成每个类别对应的“文本特征向量”# 这些向量可以直接拿来和图像特征做点积,从而完成“零样本分类”# classnames: 类别名列表,比如 ['cat', 'dog']# templates: 模板列表,比如 ["a photo of a {}", "a picture of a {}", "an image of a {}"],用于生成文本提示# model: CLIP模型,用于生成文本特征向量def zeroshot_classifier(self, classnames, templates, model):# 关闭梯度计算,因为这是推理阶段with torch.no_grad():# 初始化一个列表,用来存放每个类别的向量zeroshot_weights = []# 遍历每个类别for classname in classnames:# 根据模板生成句子,比如 ["a photo of a cat", "a picture of a cat", "an image of a cat"],用于生成文本提示# 这些句子被称为 prompt(提示语),CLIP 模型用这些文本来表示“猫”这个类的语义含义texts = [template.format(classname) for template in templates]# 把句子变成 token id(CLIP 自带的 tokenizer)texts = clip.tokenize(texts, truncate=True).cuda()# 把上面的文本 token 喂进 CLIP 的文本编码器,得到每个句子的语义向量# 输出形状是 [T, D],T 是句子数量(templates),D 是向量维度(512或768)class_embeddings = model.encode_text(texts)# 归一化每个向量,使其长度为1,归一化后,向量之间的相似度(点积)就是余弦相似度class_embeddings /= class_embeddings.norm(dim=-1, keepdim=True)# 量如果一个类别用了多个模板(比如 3 个不同句式),取平均值让它们综合起来# 这就得到这个类别的“平均语义表示”,称为“原型向量”class_embedding = class_embeddings.mean(dim=0)# 再次归一化class_embedding /= class_embedding.norm()# 把这个类别的特征(一个向量)保存到列表中zeroshot_weights.append(class_embedding)# 最终形状为 [D, C],D 是向量维度,C 是类别数量zeroshot_weights = torch.stack(zeroshot_weights, dim=1).cuda()return zeroshot_weightsdef forward(self,image: torch.Tensor = None,text = None) -> SDOutput:# 类名列表,如 ['cat','dog',...]text = self.pattern_dictionary[self.config.visual_pattern]# 关闭梯度计算,因为这是推理阶段with torch.no_grad():imagenet_templates = ['{}',]zeroshot_weights = self.zeroshot_classifier(text, imagenet_templates, self.class_model.model.float()).float()# zeroshot_weights.T 是转置,转置后形状 [num_classes, embedding_dim]self.class_model.final_fc.weight.data = zeroshot_weights.T# .contiguous() 确保张量在内存中是连续存储的,避免后面计算时出错self.class_model.final_fc.weight.data = self.class_model.final_fc.weight.data.contiguous()# 如果 text=['cat','dog','bird'] → classes=[0,1,2]classes = [i for i in range(len(text))]discrimi_image = self.resize_transform_discrimi(image)genera_image = image# real_BS:当前输入图片的批量大小(Batch Size),即一次输入多少张图real_BS = image.shape[0]after_DF_expand_BS = real_BS*self.config.input.batch_size# VAE编码图像到潜在空间self.vae, self.text_encoder, self.unet = self.vae.to(torch.float32), self.text_encoder.to(torch.float32), self.unet.to(torch.float32)# image_renormalizer:把图像像素变换到 VAE 期望的范围(比如 [0,1] → [-1,1])# .detach() 表示“断开反向传播的梯度”,让这张图像在之后的计算中不再参与梯度更新renormed_image = self.image_renormalizer(genera_image).detach()x0 = self.vae.encode(renormed_image).latent_dist.mean.float()# 乘 0.18215:Stable Diffusion(Diffusers) 的尺度常数,把 VAE 的输出缩放到 UNet 训练时的数值范围latent = x0 * 0.18215# prepare_total_timesteps 训练时的总步数(比如 1000)total_timestep = self.scheduler.num_train_timestepsfor step in range(self.config.tta.gradient_descent.train_steps):# 采样时间步 + 加噪timesteps = initiate_time_steps(step, total_timestep, after_DF_expand_BS, self.config).long()timesteps = timesteps.cuda()c, h, w = latent.shape[1:]if not self.config.tta.use_same_noise_among_timesteps:noise = torch.randn((real_BS* self.config.input.batch_size, c, h, w)).cuda()else:noise = torch.randn((1, c, h, w)).cuda()noise = noise.repeat(real_BS* self.config.input.batch_size, 1, 1, 1)if self.config.tta.adapt_topk == -1:image_features, logits, _, _, _ = self.classify(discrimi_image, classes)pred_top_idx = Noneelse:# CLIP提取特征image_features, logits, pred_top_idx, _, _ = self.classify(discrimi_image, classes)real_BS, C = logits.shape[:2]# Pick top-K predictionsif pred_top_idx is not None:pred_top_idx = pred_top_idx.squeeze(0)else:pred_top_idx = torch.arange(C).cuda()logits = logits[:, pred_top_idx]# 准备文本嵌入class_text_embeddings = prepare_class_text_embeddings(self.tokenizer, self.text_encoder, class_names=text)# 断开梯度class_text_embeddings = class_text_embeddings.detach()class_text_embeddings = class_text_embeddings[pred_top_idx, :]# Compute conditional text embeddings using weighted-summed predictions# 基于CLIP预测概率加权文本嵌入probs = logits.softmax(-1)# 把形状 [B, C] → [B, C, 1, 1]probs = probs[:, :, None, None]class_text_embeddings = (class_text_embeddings.unsqueeze(0).repeat(after_DF_expand_BS, 1, 1, 1))# 得到类别数 word_num_, word_num, _, _ = probs.shapeprobs = probs.unsqueeze(1).repeat(1,self.config.input.batch_size,1,1,1).reshape(-1,word_num,1,1)# 对每个类别的文本向量按 CLIP 预测概率加权平均# 结果 context 是每张图片的融合语义向量,包含了它“像什么”的概率信息context = (probs * class_text_embeddings).sum(1)# 把 CLIP 的视觉特征维度映射到与扩散模型 UNet 相同的特征空间image_features = self.visual_proj(image_features)# 融合视觉和文本信息context = context + image_features# pred_noise其实就是添加到图像上的噪声# 扩散模型预测噪声pred_noise = self._unet_pred_noise(x_start=latent, t=timesteps, noise=noise, context=context).float()# 计算预测误差,这个损失越小,说明模型预测的噪声越接近真实# L1:绝对误差# MSE:均方误差if self.config.tta.loss == "l1":loss = torch.nn.functional.l1_loss(pred_noise, noise)else:loss = torch.nn.functional.mse_loss(pred_noise, noise)# 如果不是最后一步,就反向传播优化CLIP特征if step != (self.config.tta.gradient_descent.train_steps-1):# 反向传播优化CLIP特征loss.backward()return SDOutput(loss=loss)
2、其他内容学习
2.1 代码实战中遇到的问题
在大模型的项目应用中几乎都要用Flash attention进行加速,但是在下载过程中很容易就会卡住,基于此我写了一篇快速安装的教程。
Flash attention快速安装
2.2 Deepseek-OCR
Deepseek-ocr最近很火,进行了学习,整理学习的笔记如下:
Vary多模态模型:DeepSeek-OCR前传?
DeepSeek-OCR:用视觉压缩文本,解决LLM长序列难题
总结
本周通过代码复现与技术学习相结合,在多模态领域取得了扎实进展:在代码实践层面,深入理解了扩散模型与CLIP的协同工作机制,掌握了VAE编码器将图像压缩至隐空间、UNet基于时间步的噪声预测、以及CLIP特征与文本嵌入的加权融合等关键技术细节;在问题解决层面,成功攻克了Flash Attention的安装难题,积累了重要的工程经验;在理论研究层面,系统梳理了DeepSeek-OCR的技术发展脉络,从Vary模型的模块化设计到视觉压缩文本解决长序列问题的创新思路。技术收获表明:代码级的深入分析是理解复杂模型工作机制的有效途径,而工程问题的解决能力同样至关重要。通过本周学习,不仅加深了对多模态模型技术细节的理解,更提升了实际问题解决能力,为后续开展更复杂的模型优化和应用开发奠定了坚实基础。
