💛前情提要💛
本文是传知代码平台中的相关前沿知识与技术的分享~
接下来我们即将进入一个全新的空间，对技术有一个全新的视角~
本文所涉及所有资源均在传知代码平台可获取
以下的内容一定会让你对AI 赋能时代有一个颠覆性的认识哦！！！
以下内容干货满满，跟上步伐吧~
📌导航小助手📌
- 💡本章重点
- 🍞一. 论文概述
- 🍞二. 扩散模型
- 🍞三.DiT架构综述
- 🍞四.代码
- 🫓总结
💡本章重点
- 生成一切的基础，DiT复现
🍞一. 论文概述
DiT的全称是Diffusion in Transformer，它将Transformer引入到扩散模型中，替换了Stable Diffusion中常用的U-Net主干。通过增加Transformer深度/宽度或增加输入令牌数量，具有较高Gflops（浮点数运算次数）的DiT始终具有较低的FID（Fréchet初始距离，用于描述生成的图片和原始的图片之间的距离）。本文会讲解扩散模型的原理，从零开始逐步复现DiT模型。
🍞二. 扩散模型
扩散模型（Diffusion Models）是一类用于生成数据（如图像、音频等）的深度学习模型，它们在近年来逐渐成为生成模型领域的重要研究方向，尤其是在图像生成任务中表现出与生成对抗网络（GANs）相媲美甚至更优的性能。
基本原理
扩散模型的基本原理是模拟数据从真实分布逐渐转变到已知分布（通常是高斯噪声分布）的过程，然后再从这个已知分布逆向恢复出数据。这个过程可以分为两个主要阶段：扩散（或正向过程）和逆扩散（或反向过程）。
扩散（正向过程）：
从原始数据开始，逐步添加随机噪声，直到数据变成不包含任何原始信息的高斯噪声。
这个过程通常是固定的，不需要学习。每一步添加的噪声量可以是一个预定义的 schedule。
扩散过程可以用一个马尔可夫链来表示，每一步都依赖于前一步的状态。
具体过程：
逆扩散（反向过程）：
从高斯噪声开始，通过学习到的模型逐步去除噪声，恢复出原始数据。
这个过程需要通过神经网络来学习，网络需要预测在每一步应该如何去除噪声，以恢复出接近原始数据的状态。
逆扩散过程通常涉及到条件概率的估计，即给定当前噪声数据，预测原始数据在该步骤的潜在状态。
🍞三.DiT架构综述
DiT的核心特点：
结合Transformer架构的扩散模型：DiT使用Transformer作为扩散模型的骨干网络，而不是传统的卷积神经网络（如U-Net）。这使得模型能够通过自注意力机制捕捉图像中的长距离依赖关系。
潜在空间操作：DiT在潜在空间中训练，通常比直接在像素空间训练更高效。通过使用变分自编码器（VAE）将图像编码到潜在空间，DiT减少了计算复杂度。
可扩展性：DiT展示了出色的可扩展性，通过增加模型的计算量（以Gflops衡量），可以显著提高生成图像的质量。这种可扩展性允许DiT在不同的分辨率和复杂度下生成图像。
DiT的工作原理：
数据准备：使用预训练的变分自编码器（VAE）将输入图像编码成潜在空间的表示。
分块化（Patchification）：将输入的潜在表示分割成一系列的小片段（patches），每个片段对应于Transformer模型的一个输入标记（token）。
Transformer Blocks模块：输入的标记序列通过一系列的Transformer块进行处理，这些块包括自注意力层、前馈神经网络以及层归一化等组件。
条件扩散过程：在训练过程中，DiT模型学习逆向扩散过程，即从噪声数据中恢复出清晰的图像。
样本生成：在训练完成后，可以通过DiT模型生成新的图像。首先，从标准正态分布中采样一个潜在表示，然后通过DiT模型逆向扩散过程，逐步去除噪声，最终解码回像素空间，得到生成的图像。
🍞四.代码
前向加噪
import torch
from config import *# 前向diffusion计算参数
betas=torch.linspace(0.0001,0.02,T) # (T,)
alphas=1-betas # (T,)
alphas_cumprod=torch.cumprod(alphas,dim=-1) # alpha_t累乘 (T,) [a1,a2,a3,....] -> [a1,a1*a2,a1*a2*a3,.....]
alphas_cumprod_prev=torch.cat((torch.tensor([1.0]),alphas_cumprod[:-1]),dim=-1) # alpha_t-1累乘 (T,), [1,a1,a1*a2,a1*a2*a3,.....]
variance=(1-alphas)*(1-alphas_cumprod_prev)/(1-alphas_cumprod) # denoise用的方差 (T,)# 执行前向加噪
def forward_add_noise(x,t): # batch_x: (batch,channel,height,width), batch_t: (batch_size,)noise=torch.randn_like(x) # 为每张图片生成第t步的高斯噪音 (batch,channel,height,width)batch_alphas_cumprod=alphas_cumprod[t].view(x.size(0),1,1,1) x=torch.sqrt(batch_alphas_cumprod)*x+torch.sqrt(1-batch_alphas_cumprod)*noise # 基于公式直接生成第t步加噪后图片return x,noise
DiT模型
- 第一步：定义PatchEmbedding
class PatchEmbedding(nn.Module):def __init__(self,channel,patch_size,emb_size):super().__init__()self.patch_size=patch_sizeself.patch_count=img_size//self.patch_sizeself.channel=channelself.conv=nn.Conv2d(in_channels=channel,out_channels=channel*patch_size**2,kernel_size=patch_size,padding=0,stride=patch_size) self.patch_emb=nn.Linear(in_features=channel*patch_size**2,out_features=emb_size) self.patch_pos_emb=nn.Parameter(torch.rand(1,self.patch_count**2,emb_size))def forward(self,x):x=self.conv(x) # (batch,new_channel,patch_count,patch_count)x=x.permute(0,2,3,1) # (batch,patch_count,patch_count,new_channel)x=x.view(x.size(0),self.patch_count*self.patch_count,x.size(3)) # (batch,patch_count**2,new_channel)x=self.patch_emb(x) # (batch,patch_count**2,emb_size)x=x+self.patch_pos_emb # (batch,patch_count**2,emb_size)return x
- 第二步：定义time_embedding
class TimeEmbedding(nn.Module):def __init__(self,emb_size):super().__init__()self.half_emb_size=emb_size//2half_emb=torch.exp(torch.arange(self.half_emb_size)*(-1*math.log(10000)/(self.half_emb_size-1)))self.register_buffer('half_emb',half_emb)def forward(self,t):t=t.view(t.size(0),1)half_emb=self.half_emb.unsqueeze(0).expand(t.size(0),self.half_emb_size)half_emb_t=half_emb*tembs_t=torch.cat((half_emb_t.sin(),half_emb_t.cos()),dim=-1)return embs_t
- 第三步：定义DiT Block
class DiTBlock(nn.Module):def __init__(self,emb_size,nhead=4):super().__init__()# 定义各种参数，包括线形层，Transformerdef forward(self,x,condition):# 计算
- 第四步：定义DiT
class DiT(nn.Module):def __init__(self,img_size,patch_size,channel,emb_size,label_num,dit_num,head):super().__init__()# 定义标签的Embedding，dit_num个DiT block，Layernorm等def forward(self,x,label):# 计算
- 开始训练
for epoch in range(EPOCH):for imgs,labels in dataloader:x=imgs*2-1 # 图像的像素范围从[0,1]转换到[-1,1],和噪音高斯分布范围对应t=torch.randint(0,T,(imgs.size(0),)) # 为每张图片生成随机t时刻y=labelsx,noise=forward_add_noise(x,t) # x:加噪图 noise:噪音pred_noise=model(x.to(DEVICE),t.to(DEVICE),y.to(DEVICE))loss=loss_fn(pred_noise,noise.to(DEVICE))optimzer.zero_grad()loss.backward()optimzer.step()
🫓总结
综上，我们基本了解了“一项全新的技术啦” 🍭 ~~
恭喜你的内功又双叒叕得到了提高！！！
感谢你们的阅读😆
后续还会继续更新💓，欢迎持续关注📌哟~
💫如果有错误❌，欢迎指正呀💫
✨如果觉得收获满满，可以点点赞👍支持一下哟~✨
【传知科技 – 了解更多新知识】