手动实现一个扩散模型DDPM

!pip install -q -U einops datasets matplotlib tqdm
import mathfrom inspect import isfunctionfrom functools import partial                                 
%matplotlib inlineimport matplotlib.pyplot as pltfrom tqdm.auto import tqdmfrom einops import rearrange, reducefrom einops.layers.torch import Rearrange
import torchfrom torch import nn, einsumimport torch.nn.functional as F

import torch
# cosine_beta_schedule函数用于创建一个余弦退火beta调度。# 这种调度方法基于余弦函数，并且可以调整随时间的衰减速率。def cosine_beta_schedule(timesteps, s=0.008):    steps = timesteps + 1  # 计算总的步数，需要比时间步多一个，以便计算alpha的累积乘积    x = torch.linspace(0, timesteps, steps)  # 创建从0到timesteps的均匀分布的张量    # 计算alpha的累积乘积，使用一个余弦变换，并平方来计算当前步的alpha值    alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * torch.pi * 0.5) ** 2    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]  # 归一化，确保初始值为1    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])  # 计算每个时间步的beta值    return torch.clip(betas, 0.0001, 0.9999)  # 对beta值进行裁剪，避免过大或过小
# linear_beta_schedule函数用于创建一个线性退火beta调度。# 这意味着beta值将从beta_start线性增加到beta_end。def linear_beta_schedule(timesteps):    beta_start = 0.0001  # 定义起始beta值    beta_end = 0.02  # 定义结束beta值    return torch.linspace(beta_start, beta_end, timesteps)  # 创建一个线性分布的beta值数组
# quadratic_beta_schedule函数用于创建一个二次退火beta调度。# 这意味着beta值将根据二次函数变化。def quadratic_beta_schedule(timesteps):    beta_start = 0.0001  # 定义起始beta值    beta_end = 0.02  # 定义结束beta值    # 创建一个线性分布的数组，然后将其平方以生成二次分布，最后再次平方以计算beta值    return torch.linspace(beta_start**0.5, beta_end**0.5, timesteps) ** 2
# sigmoid_beta_schedule函数用于创建一个sigmoid退火beta调度。# 这意味着beta值将根据sigmoid函数变化，这是一种常见的激活函数。def sigmoid_beta_schedule(timesteps):    beta_start = 0.0001  # 定义起始beta值    beta_end = 0.02  # 定义结束beta值    betas = torch.linspace(-6, 6, timesteps)  # 创建一个从-6到6的线性分布，用于sigmoid函数的输入    # 应用sigmoid函数，并根据beta_start和beta_end调整其范围和位置    return torch.sigmoid(betas) * (beta_end - beta_start) + beta_start

# import torch  # 假设在代码的其他部分已经导入了torch库
# 定义前向扩散函数# x_start: 初始数据，例如一批图像# t: 扩散的时间步，表示当前的扩散阶段# noise: 可选参数，如果提供，则使用该噪声数据；否则，将生成新的随机噪声def q_sample(x_start, t, noise=None):    if noise is None:        noise = torch.randn_like(x_start)  # 如果未提供噪声，则生成一个与x_start形状相同的随机噪声张量
    # 提取对应于时间步t的α的累积乘积的平方根    sqrt_alphas_cumprod_t = extract(sqrt_alphas_cumprod, t, x_start.shape)    # 提取对应于时间步t的1-α的累积乘积的平方根    sqrt_one_minus_alphas_cumprod_t = extract(        sqrt_one_minus_alphas_cumprod, t, x_start.shape    )
    # 返回前向扩散的结果，该结果是初始数据和噪声的线性组合    # 系数sqrt_alphas_cumprod_t和sqrt_one_minus_alphas_cumprod_t分别用于缩放初始数据和噪声    return sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise

# take time stepfor noise in [10,20,40,80 100]:  t = torch.tensor([40])  get_noisy_image(x_start, t)

import torch.nn as nnfrom einops import rearrange  # 假设已经导入了einops库中的rearrange函数from torch_utils import exists  # 假设已经定义了exists函数，用于检查对象是否存在
# 定义一个基础的Block类，该类将作为神经网络中的一个基本构建模块class Block(nn.Module):    def __init__(self, dim, dim_out, groups=8):        super().__init__()        # 一个2D卷积层，卷积核大小为3x3，边缘填充为1，从输入维度dim到输出维度dim_out        self.proj = nn.Conv2d(dim, dim_out, 3, padding=1)        # GroupNorm层用于归一化，分组数为groups        self.norm = nn.GroupNorm(groups, dim_out)        # 使用SiLU（也称为Swish）作为激活函数        self.act = nn.SiLU()            def forward(self, x, scale_shift=None):        x = self.proj(x)  # 应用卷积操作        x = self.norm(x)  # 应用归一化操作        # 如果scale_shift参数存在，则对归一化后的数据进行缩放和位移操作        if exists(scale_shift):            scale, shift = scale_shift            x = x * (scale + 1) + shift        x = self.act(x)  # 应用激活函数        return x  # 返回处理后的数据
# 定义一个ResnetBlock类，用于构建残差网络中的基本块class ResnetBlock(nn.Module):    """https://arxiv.org/abs/1512.03385"""        def __init__(self, dim, dim_out, *, time_emb_dim=None, groups=8):        super().__init__()        # 如果time_emb_dim存在，定义一个小型的多层感知器（MLP）网络        self.mlp = (            nn.Sequential(nn.SiLU(), nn.Linear(time_emb_dim, dim_out))            if exists(time_emb_dim)            else None        )        # 定义两个顺序的基础Block模块        self.block1 = Block(dim, dim_out, groups=groups)        self.block2 = Block(dim_out, dim_out, groups=groups)        # 如果输入维度dim和输出维度dim_out不同，则使用1x1卷积进行维度调整        # 否则使用Identity层（相当于不做任何处理）        self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()            def forward(self, x, time_emb=None):        h = self.block1(x)  # 通过第一个Block模块        # 如果存在时间嵌入向量time_emb且存在mlp模块，则将其应用到h上        if exists(self.mlp) and exists(time_emb):            time_emb = self.mlp(time_emb)  # 通过MLP网络            # 重整time_emb的形状以匹配h的形状，并将结果加到h上            h = rearrange(time_emb, "b c -> b c 1 1") + h        h = self.block2(h)  # 通过第二个Block模块        return h + self.res_conv(x)  # 将Block模块的输出与调整维度后的原始输入x相加并返回

import torchfrom torch import nnfrom einops import rearrangeimport torch.nn.functional as F
# 定义一个标准的多头注意力（Multi-Head Attention）机制的类class Attention(nn.Module):    def __init__(self, dim, heads=4, dim_head=32):        super().__init__()        # 根据维度的倒数平方根来缩放查询（Query）向量        self.scale = dim_head ** -0.5        # 头的数量（多头中的"多"）        self.heads = heads        # 计算用于多头注意力的隐藏层维度        hidden_dim = dim_head * heads        # 定义一个卷积层将输入的特征映射到QKV（查询、键、值）空间        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)        # 定义一个卷积层将多头注意力的输出映射回原特征空间        self.to_out = nn.Conv2d(hidden_dim, dim, 1)
    def forward(self, x):        # 获取输入的批量大小、通道数、高度和宽度        b, c, h, w = x.shape        # 使用to_qkv卷积层得到QKV，并将其分离为三个组件        qkv = self.to_qkv(x).chunk(3, dim=1)        # 将QKV重排并缩放查询向量        q, k, v = map(            lambda t: rearrange(t, "b (h c) x y -> b h c (x y)", h=self.heads), qkv        )        q = q * self.scale        # 使用爱因斯坦求和约定计算查询和键之间的相似度得分        sim = einsum("b h d i, b h d j -> b h i j", q, k)        # 从相似度得分中减去最大值以提高数值稳定性        sim = sim - sim.amax(dim=-1, keepdim=True).detach()        # 应用Softmax函数获取注意力权重        attn = sim.softmax(dim=-1)        # 使用注意力权重对值进行加权        out = einsum("b h i j, b h d j -> b h i d", attn, v)        # 将输出重新排列回原始的空间形状        out = rearrange(out, "b h (x y) d -> b (h d) x y", x=h, y=w)        # 返回通过输出卷积层的结果        return self.to_out(out)
# 定义一个线性注意力（Linear Attention）机制的类class LinearAttention(nn.Module):    def __init__(self, dim, heads=4, dim_head=32):        super().__init__()        # 根据维度的倒数平方根来缩放查询（Query）向量        self.scale = dim_head ** -0.5        # 头的数量        self.heads = heads        # 计算用于多头注意力的隐藏层维度        hidden_dim = dim_head * heads        # 定义一个卷积层将输入的特征映射到QKV空间        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)        # 定义一个顺序容器包含卷积层和组归一化层将输出映射回原特征空间        self.to_out = nn.Sequential(nn.Conv2d(hidden_dim, dim, 1),                                    nn.GroupNorm(1, dim))
    def forward(self, x):        # 获取输入的批量大小、通道数、高度和宽度        b, c, h, w = x.shape        # 使用to_qkv卷积层得到QKV，并将其分离为三个组件        qkv = self.to_qkv(x).chunk(3, dim=1)        # 将QKV重排，应用Softmax函数并缩放查询向量        q, k, v = map(            lambda t: rearrange(t, "b (h c) x y -> b h c (x y)", h=self.heads), qkv        )        q = q.softmax(dim=-2)        k = k.softmax(dim=-1)        q = q * self.scale        # 计算上下文矩阵，是键和值的加权组合        context = torch.einsum("b h d n, b h e n -> b h d e", k, v)        # 使用上下文矩阵和查询计算最终的注意力输出        out = torch.einsum("b h d e, b h d n -> b h e n", context, q)        # 将输出重新排列回原始的空间形状        out = rearrange(out, "b h c (x y) -> b (h c) x y", h=self.heads, x=h, y=w)        # 返回经过输出顺序容器处理的结果        return self.to_out(out)

class SinusoidalPositionEmbeddings(nn.Module):    def __init__(self, dim):        super().__init__()        self.dim = dim
    def forward(self, time):        device = time.device        half_dim = self.dim // 2        embeddings = math.log(10000) / (half_dim - 1)        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)        embeddings = time[:, None] * embeddings[None, :]        embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)        return embeddings

class Unet(nn.Module):    # 初始化函数，定义U-Net网络的结构和参数    def __init__(            self,            dim,  # 基本隐藏层维度            init_dim=None,  # 初始层维度，如果未提供则会根据dim计算得出            out_dim=None,  # 输出维度，如果未提供则默认为输入图像的通道数            dim_mults=(1, 2, 4, 8),  # 控制每个阶段隐藏层维度倍增的倍数            channels=3,  # 输入图像的通道数，默认为3            with_time_emb=True,  # 是否使用时间嵌入，这对于某些生成模型可能是必要的            resnet_block_groups=8,  # ResNet块中的组数            use_convnext=True,  # 是否使用ConvNeXt块而不是ResNet块            convnext_mult=2,  # ConvNeXt块的维度倍增因子    ):        super().__init__()  # 调用父类构造函数        # 确定各层维度        self.channels = channels        init_dim = default(init_dim, dim // 3 * 2)  # 设置或计算初始层维度        self.init_conv = nn.Conv2d(channels, init_dim, 7, padding=3)  # 初始卷积层，使用7x7卷积核和padding        dims = [init_dim, *map(lambda m: dim * m, dim_mults)]  # 计算每个阶段的维度        in_out = list(zip(dims[:-1], dims[1:]))  # 创建输入输出维度对        # 根据use_convnext选择块类        if use_convnext:            block_klass = partial(ConvNextBlock, mult=convnext_mult)        else:            block_klass = partial(ResnetBlock, groups=resnet_block_groups)        # 时间嵌入层        if with_time_emb:            time_dim = dim * 4  # 时间嵌入的维度            self.time_mlp = nn.Sequential(  # 时间嵌入的多层感知机                SinusoidalPositionEmbeddings(dim),  # 正弦位置嵌入                nn.Linear(dim, time_dim),  # 线性变换                nn.GELU(),  # GELU激活函数                nn.Linear(time_dim, time_dim),  # 再一次线性变换            )        else:            time_dim = None            self.time_mlp = None        # 下采样层        self.downs = nn.ModuleList([])        self.ups = nn.ModuleList([])        num_resolutions = len(in_out)  # 解析的层数        # 构建下采样模块        for ind, (dim_in, dim_out) in enumerate(in_out):            is_last = ind >= (num_resolutions - 1)  # 是否为最后一层            self.downs.append(  # 添加下采样块                nn.ModuleList(                    [                        block_klass(dim_in, dim_out, time_emb_dim=time_dim),  # 卷积块                        block_klass(dim_out, dim_out, time_emb_dim=time_dim),  # 卷积块                        Residual(PreNorm(dim_out, LinearAttention(dim_out))),  # 残差连接和注意力模块                        Downsample(dim_out) if not is_last else nn.Identity(),  # 下采样或恒等映射                    ]                )            )        # 中间层（瓶颈层）        mid_dim = dims[-1]        # 中间层（瓶颈层）        # 第一个中间卷积块        self.mid_block1 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim)        # 中间层的注意力模块        self.mid_attn = Residual(PreNorm(mid_dim, Attention(mid_dim)))        # 第二个中间卷积块        self.mid_block2 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim)
        # 构建上采样模块        for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):            is_last = ind >= (num_resolutions - 1)  # 是否是最后一次上采样，减2是因为我们需要留出一个输出层            self.ups.append(                nn.ModuleList(                    [                        # 卷积块，这里输入维度翻倍是因为上采样过程中会与编码器阶段的相应层进行拼接                        block_klass(dim_out * 2, dim_in, time_emb_dim=time_dim),                        # 卷积块                        block_klass(dim_in, dim_in, time_emb_dim=time_dim),                        # 残差和注意力模块                        Residual(PreNorm(dim_in, LinearAttention(dim_in))),                        # 上采样或恒等映射                        Upsample(dim_in) if not is_last else nn.Identity(),                    ]                )            )
        # 设置或计算输出维度，如果未提供则默认为输入图像的通道数        out_dim = default(out_dim, channels)        # 最后的卷积层，将输出维度变换到期望的输出维度        self.final_conv = nn.Sequential(            block_klass(dim, dim),  # 卷积块            nn.Conv2d(dim, out_dim, 1)  # 1x1卷积，用于输出维度变换        )
    # 前向传播函数    def forward(self, x, time):        # 初始卷积层        x = self.init_conv(x)        # 如果存在时间嵌入层，则将时间编码        t = self.time_mlp(time) if exists(self.time_mlp) else None        # 用于存储各个阶段的特征图        h = []
        # 下采样过程        for block1, block2, attn, downsample in self.downs:            x = block1(x, t)  # 应用卷积块            x = block2(x, t)  # 应用卷积块            x = attn(x)  # 应用注意力模块            h.append(x)  # 存储特征图以便后续的拼接            x = downsample(x)  # 应用下采样或恒等映射
        # 中间层或瓶颈层        x = self.mid_block1(x, t)  # 第一个中间卷积块        x = self.mid_attn(x)  # 中间层的注意力模块        x = self.mid_block2(x, t)  # 第二个中间卷积块
        # 上采样过程        for block1, block2, attn, upsample in self.ups:            # 拼接特征图和对应的编码器阶段的特征图            x = torch.cat((x, h.pop()), dim=1)            x = block1(x, t)  # 应用卷积块            x = block2(x, t)  # 应用卷积块            x = attn(x)  # 应用注意力模块            x = upsample(x)  # 应用上采样或恒等映射
        # 最后的输出层，输出最终的特征图或图像        return self.final_conv(x)

import torchimport torch.nn.functional as F
# 定义损失函数，它评估去噪模型的性能def p_losses(denoise_model, x_start, t, noise=None, loss_type="l1"):    if noise is None:        noise = torch.randn_like(x_start)  # 如果未提供噪声，则生成一个与x_start形状相同的随机噪声张量
    # 使用q_sample函数生成带有噪声的数据x_noisy，这模拟了扩散模型的前向过程    x_noisy = q_sample(x_start=x_start, t=t, noise=noise)    # 使用去噪模型对噪声数据x_noisy进行预测，试图恢复加入的噪声    predicted_noise = denoise_model(x_noisy, t)
    # 根据指定的损失类型计算损失    if loss_type == 'l1':  # 如果损失类型为L1损失        loss = F.l1_loss(noise, predicted_noise)  # 使用L1损失函数计算真实噪声和预测噪声之间的差异    elif loss_type == 'l2':  # 如果损失类型为L2损失（均方误差损失）        loss = F.mse_loss(noise, predicted_noise)  # 使用均方误差损失函数计算真实噪声和预测噪声之间的差异    elif loss_type == "huber":  # 如果损失类型为Huber损失        loss = F.smooth_l1_loss(noise, predicted_noise)  # 使用Huber损失函数，这是L1和L2损失的结合，对异常值不那么敏感    else:        raise NotImplementedError()  # 如果指定了未实现的损失类型，则抛出异常
    return loss  # 返回计算得到的损失值

if __name__=="__main__":    for epoch in range(epochs):        for step, batch in tqdm(enumerate(dataloader), desc='Training'):          optimizer.zero_grad()          batch = batch[0]
          batch_size = batch.shape[0]          batch = batch.to(device)          # 国内版启用这段，注释上面两行          # batch_size = batch[0].shape[0]          # batch = batch[0].to(device)
          # Algorithm 1 line 3: sample t uniformally for every example in the batch          t = torch.randint(0, timesteps, (batch_size,), device=device).long()
          loss = p_losses(model, batch, t, loss_type="huber")
          if step % 50 == 0:            print("Loss:", loss.item())
          loss.backward()          optimizer.step()
          # save generated images          if step != 0 and step % save_and_sample_every == 0:            milestone = step // save_and_sample_every            batches = num_to_groups(4, batch_size)            all_images_list = list(map(lambda n: sample(model, batch_size=n, channels=channels), batches))            all_images = torch.cat(all_images_list, dim=0)            all_images = (all_images + 1) * 0.5            # save_image(all_images, str(results_folder / f'sample-{milestone}.png'), nrow = 6)            currentDateAndTime = datetime.now()            torch.save(model,f"train.pt")