免费网站源码html,设计制作我们的小船,wordpress回帖可见,海口建站模板系统文章目录 序言一参考资料①代码来源②相关概念理解③公式推导及训练流程讲解④搜索问题的网站⑤模型运行的环境 二代码解读①模型②训练③测试 三主要训练过程的解析 序言
本文主要对一个基于minist数据集搭建的DDPM模型代码中各个模块的含义进行解析#xff0c;初步记录了自… 文章目录 序言一参考资料①代码来源②相关概念理解③公式推导及训练流程讲解④搜索问题的网站⑤模型运行的环境 二代码解读①模型②训练③测试 三主要训练过程的解析 序言
本文主要对一个基于minist数据集搭建的DDPM模型代码中各个模块的含义进行解析初步记录了自己了解扩散模型的一个过程为后续的进一步学习打基础。文中的错误之处还望大家批评指正一参考资料
①代码来源
参考的代码来源
②相关概念理解
超详细的扩散模型Diffusion Models原理代码
③公式推导及训练流程讲解
DDPM1 DDPM2
④搜索问题的网站
geekgpt 此网站可以对代码进行注释对公式推导的流程进行解释利用的好可以帮助我们更好的理解我们所的遇到的大部分问题
⑤模型运行的环境
本文是在google drive的colab中运行的 colab的部署过程可以参考以下内容Colab 实用教程
二代码解读
①模型
import os
import math
from abc import abstractmethodimport numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdmdef timestep_embedding(timesteps, dim, max_period10000):Create sinusoidal timestep embeddings.Args:timesteps (Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional.dim (int): the dimension of the output.max_period (int, optional): controls the minimum frequency of the embeddings. Defaults to 10000.Returns:Tensor: an [N x dim] Tensor of positional embeddings.# 计算嵌入向量的一半维度half dim // 2# 计算频率用来生成正弦和余弦成分freqs torch.exp(-math.log(max_period) * torch.arange(start0, endhalf, dtypetorch.float32) / half).to(devicetimesteps.device)# 计算角度参数用于生成正弦和余弦成分#进行维度的扩充将1*step转化为step*1 然后和1*half进行矩阵运算将数据的维度扩充到了half维度half为偶数args timesteps[:, None].float() * freqs[None]# 生成正弦和余弦成分然后连接它们以形成嵌入向量embedding torch.cat([torch.cos(args), torch.sin(args)], dim-1)# 如果维度是奇数添加一个额外的零维度if dim % 2:embedding torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim-1)# 返回时间步嵌入向量return embeddingclass TimestepBlock(nn.Module):Any module where forward() takes timestep embeddings as a second argument.abstractmethoddef forward(self, x, t):Apply the module to x given t timestep embeddings.passclass TimestepEmbedSequential(nn.Sequential, TimestepBlock):A sequential module that passes timestep embeddings to the children that support it as an extra input.def forward(self, x, t):for layer in self:if isinstance(layer, TimestepBlock):x layer(x, t)else:x layer(x)return x# layer 是 TimestepEmbedSequential 类中的每个子模块layer的引用。在这个循环中我们遍历了 TimestepEmbedSequential 类中的每个子模块并对其进行操作。如果子模块是 TimestepBlock 的实例则调用其 forward() 方法并将输入数据 x 和时间步骤嵌入向量 t 传递给它否则我们只是将输入数据 x 传递给子模块#参数channels指定了归一化层的通道数而nn.GroupNorm的第一个参数32表示将输入数据的通道分成32个子组每个子组内的特征将被独立地归一化组归一化的主要作用是解决深度神经网络中的内部协变量偏移问题提高模型的训练稳定性使其更适合处理不同批量大小和高分辨率数据同时也有助于模型的泛化能力。
def norm_layer(channels):return nn.GroupNorm(32, channels)class ResidualBlock(TimestepBlock):def __init__(self, in_channels, out_channels, time_channels, dropout):super().__init__()self.conv1 nn.Sequential(norm_layer(in_channels),nn.SiLU(),nn.Conv2d(in_channels, out_channels, kernel_size3, padding1))# pojection for time step embeddingself.time_emb nn.Sequential(nn.SiLU(),nn.Linear(time_channels, out_channels))self.conv2 nn.Sequential(norm_layer(out_channels),nn.SiLU(),nn.Dropout(pdropout),nn.Conv2d(out_channels, out_channels, kernel_size3, padding1))if in_channels ! out_channels:self.shortcut nn.Conv2d(in_channels, out_channels, kernel_size1)else:self.shortcut nn.Identity()def forward(self, x, t):x has shape [batch_size, in_dim, height, width]t has shape [batch_size, time_dim]h self.conv1(x)# Add time step embeddingsh self.time_emb(t)[:, :, None, None]h self.conv2(h)return h self.shortcut(x)class AttentionBlock(nn.Module):def __init__(self, channels, num_heads1):Attention block with shortcutArgs:channels (int): channelsnum_heads (int, optional): attention heads. Defaults to 1.super().__init__()self.num_heads num_headsassert channels % num_heads 0self.norm norm_layer(channels)self.qkv nn.Conv2d(channels, channels * 3, kernel_size1, biasFalse)self.proj nn.Conv2d(channels, channels, kernel_size1)def forward(self, x):B, C, H, W x.shapeqkv self.qkv(self.norm(x))#将模型的维度扩充3倍q, k, v qkv.reshape(B*self.num_heads, -1, H*W).chunk(3, dim1)scale 1. / math.sqrt(math.sqrt(C // self.num_heads))# 计算了一个用于缩放注意力分数的标度因子scaling factor。这个标度因子通常用于控制注意力分数的大小以避免过大的数值有助于稳定训练过程attn torch.einsum(bct,bcs-bts, q * scale, k * scale)#这一行代码执行了一个张量乘法操作并计算了注意力分数attention scoresattn attn.softmax(dim-1)#进行 softmax 归一化以确保每个位置的分数都在 [0, 1] 范围内h torch.einsum(bts,bcs-bct, attn, v)#torch.einsum 函数是一个强大的张量运算工具它允许用户根据一种命名约定来指定张量的操作以实现高效的张量操作和组合。它的基本语法是#result torch.einsum(ij,jk-ik, A, B)#这将计算两个矩阵 A 和 B 的矩阵乘法。字符串 ij,jk-ik 描述了两个矩阵的操作其中 ij 表示 A 的行和列而 jk 表示 B 的行和列最终得到一个矩阵#其中 attn 和 v 是输入张量具有以下维度#attn 的形状为 (batch_size, sequence_length, num_heads)#其中 batch_size 表示批处理大小sequence_length 表示序列长度num_heads 表示注意力头的数量。#v 的形状为 (batch_size, sequence_length, value_dim)其中 value_dim 表示每个注意力值的维度。#输出张量的形状为 (batch_size, sequence_length, value_dim)它表示了加权和的结果其中每个元素#都是通过将 attn 中的权重应用到 v 中的相应部分来计算的。#这种操作通常用于多头注意力机制中其中 attn 包含了注意力分数或权重v 包含了值而输出则是根#据权重对值进行加权求和的结果。这有助于模型在自注意力机制中将不同的信息聚合到输出中。
#用于执行多头注意力操作将注意力权重应用于值并计算加权h h.reshape(B, -1, H, W)h self.proj(h)return h x# 定义一个名为Upsample的自定义神经网络模块
class Upsample(nn.Module):def __init__(self, channels, use_conv):super().__init__()# 初始化函数接受两个参数channels表示输入通道数use_conv表示是否使用卷积层# 将use_conv标记存储在模块中以便后续的操作可以根据该标记来选择不同的处理方式self.use_conv use_conv # 如果use_conv为True即选择使用卷积层if use_conv:# 创建一个卷积层输入通道数和输出通道数都为channels# 使用3x3的卷积核kernel_size3并在输入周围填充1个像素padding1self.conv nn.Conv2d(channels, channels, kernel_size3, padding1)# 定义模块的前向传播函数接受输入张量x作为参数def forward(self, x):# 使用F.interpolate函数对输入张量x进行上采样# 上采样的尺度因子为2scale_factor2将原图像的每个维度放大2倍采用最近邻插值方式modenearestx F.interpolate(x, scale_factor2, modenearest)# 如果use_conv为True即选择使用卷积层if self.use_conv:# 将上采样后的张量x输入到卷积层self.conv中进行卷积操作x self.conv(x)# 返回处理后的张量x作为模块的输出return xclass Downsample(nn.Module):
#上采样和下采样的初始化都是输入通道数和是否用卷积如果不用卷积那么就用池化层进行下采样如果用卷积那么就用卷积核步长为2去达到下采样的效果def __init__(self, channels, use_conv):super().__init__()self.use_conv use_convif use_conv:self.op nn.Conv2d(channels, channels, kernel_size3, stride2, padding1)else:#利用平均池化层将数据缩小为原来的1/2self.op nn.AvgPool2d(stride2)def forward(self, x):return self.op(x)class UNetModel(nn.Module):The full UNet model with attention and timestep embeddingdef __init__(self,in_channels3, # 输入通道数默认为3适用于RGB图像model_channels128, # 模型通道数默认为128out_channels3, # 输出通道数默认为3适用于RGB图像num_res_blocks2, # 残差块的数量默认为2attention_resolutions(8, 16), # 注意力分辨率的元组默认为(8, 16)dropout0, # Dropout概率默认为0不使用Dropoutchannel_mult(1, 2, 2, 2), # 通道倍增因子的元组默认为(1, 2, 2, 2)conv_resampleTrue, # 是否使用卷积重采样默认为Truenum_heads4 # 注意力头的数量默认为4):super().__init__()# 初始化模型的各种参数self.in_channels in_channelsself.model_channels model_channelsself.out_channels out_channelsself.num_res_blocks num_res_blocksself.attention_resolutions attention_resolutionsself.dropout dropoutself.channel_mult channel_multself.conv_resample conv_resampleself.num_heads num_heads# 时间嵌入用于处理时间信息的嵌入time_embed_dim model_channels * 4self.time_embed nn.Sequential(nn.Linear(model_channels, time_embed_dim),nn.SiLU(),nn.Linear(time_embed_dim, time_embed_dim),)# 下采样块#所有的模块都是先定义然后通过迭代的方式往模块里面加东西self.down_blocks nn.ModuleList([TimestepEmbedSequential(nn.Conv2d(in_channels, model_channels, kernel_size3, padding1))])down_block_chans [model_channels] # 存储下采样块每一阶段的通道数ch model_channels # 当前通道数初始化为模型通道数 初始为128ds 1 # 下采样的倍数初始值为1# 遍历不同阶段的下采样块#channel_mult模块为1222下采样块每层的块数for level, mult in enumerate(channel_mult):#num_res_blocks为残差块的数量表示每块需要的残差快的数量for _ in range(num_res_blocks):layers [#ch为输入通道数mult * model_channels为需要输出的维度数time_embed_dim为时间嵌入的维度ResidualBlock(ch, mult * model_channels, time_embed_dim, dropout)#初始化剩余块让我们后续能用forward函数将时间嵌入到x中]ch mult * model_channels#ds为一个值一开始为1然后每次乘以2这里如果ds为8或者16时需要加上一个注意力模块if ds in attention_resolutions:layers.append(AttentionBlock(ch, num_headsnum_heads))#将加入了残差快和注意力块的层加入下采样块当中self.down_blocks.append(TimestepEmbedSequential(*layers))#记录每一层采样的通道数down_block_chans.append(ch)if level ! len(channel_mult) - 1: # 最后一个阶段不使用下采样#这里由于之前的ch*2 所以下采样后又恢复到了 ch所以我们在下采样通道中加入的chself.down_blocks.append(TimestepEmbedSequential(Downsample(ch, conv_resample)))down_block_chans.append(ch)ds * 2#整个流程的格式变换12812864256,256128256,256# 中间块#中间块就是一个残差块注意力块残差块self.middle_block TimestepEmbedSequential(ResidualBlock(ch, ch, time_embed_dim, dropout),AttentionBlock(ch, num_headsnum_heads),ResidualBlock(ch, ch, time_embed_dim, dropout))# 上采样块self.up_blocks nn.ModuleList([])#反过来计算通道的情况2,2,21for level, mult in list(enumerate(channel_mult))[::-1]:#反向时残差块的数目为3for i in range(num_res_blocks 1):layers [ResidualBlock(ch down_block_chans.pop(),model_channels * mult,time_embed_dim,dropout)]ch model_channels * multif ds in attention_resolutions:layers.append(AttentionBlock(ch, num_headsnum_heads))#如果level不为0并且i为2时最后一块时,进行上采样if level and i num_res_blocks:layers.append(Upsample(ch, conv_resample))ds // 2self.up_blocks.append(TimestepEmbedSequential(*layers))# 输出层#只是一个正则化激活后的再一次不改变通道数的卷积self.out nn.Sequential(norm_layer(ch),nn.SiLU(),nn.Conv2d(model_channels, out_channels, kernel_size3, padding1),)def forward(self, x, timesteps):Apply the model to an input batch.Args:x (Tensor): [N x C x H x W]timesteps (Tensor): a 1-D batch of timesteps.Returns:Tensor: [N x C x ...]#记录每次下采样得到结果用于后面上采样的copy and crophs []# 时间步嵌入#利用timesteps参数计算时间步的嵌入#首先用timestep_embedding,将时间序列timesteps1*n转化为n*model_channels#然后用time_embed将之前的n*model_channels转化为 n*time_embed_dim也就是原来的mocel_channels*4emb self.time_embed(timestep_embedding(timesteps, self.model_channels))#最终得到一个时间步嵌入的矩阵# 下采样阶段h xfor module in self.down_blocks:#每次用时间步嵌入的矩阵信息emb更新并记录每次的hh module(h, emb)hs.append(h)# 中间阶段h self.middle_block(h, emb)# 上采样阶段for module in self.up_blocks:cat_in torch.cat([h, hs.pop()], dim1)h module(cat_in, emb)return self.out(h)#线性β只是等距的值
def linear_beta_schedule(timesteps):beta schedulescale 1000 / timestepsbeta_start scale * 0.0001beta_end scale * 0.02#等距生成timesteps个数值作为β的取值return torch.linspace(beta_start, beta_end, timesteps, dtypetorch.float64)#实现了一个余弦学习率调度
#timesteps: 这是一个整数参数指定生成渐变序列的时间步数。
#s: 这是余弦调度的一个超参数控制余弦曲线的形状。默认值为0.008。
def cosine_beta_schedule(timesteps, s0.008):cosine scheduleas proposed in https://arxiv.org/abs/2102.09672steps timesteps 1x torch.linspace(0, timesteps, steps, dtypetorch.float64)#alphas_cumprod: 这个步骤计算了一个余弦曲线的累积乘积并且通过缩放将其限制在0到1之间。这个曲线的形状由s参数控制。alphas_cumprod torch.cos(((x / timesteps) s) / (1 s) * math.pi * 0.5) ** 2alphas_cumprod alphas_cumprod / alphas_cumprod[0]#betas: 计算了渐变的beta值序列通过计算相邻时间步的alphas_cumprod之间的差异。betas 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])#最后将beta值序列裁剪到区间[0, 0.999]之间以确保其在有效范围内。return torch.clip(betas, 0, 0.999)class GaussianDiffusion:def __init__(self,timesteps1000, # 初始化函数设置默认时间步数为1000beta_schedulelinear # 初始化函数设置默认的beta调度为linear):self.timesteps timesteps # 存储时间步数# 根据选择的beta调度类型生成beta值的序列if beta_schedule linear:betas linear_beta_schedule(timesteps)elif beta_schedule cosine:betas cosine_beta_schedule(timesteps)else:raise ValueError(funknown beta schedule {beta_schedule})self.betas betas # 存储beta值序列# 计算alpha值1 - beta和alpha的累积乘积1,23变为1,26self.alphas 1. - self.betasself.alphas_cumprod torch.cumprod(self.alphas, axis0)#F.pada,b,c函数在a向量的最前面和最后面分别添加b个c元素self.alphas_cumprod_prev F.pad(self.alphas_cumprod[:-1], (1, 0), value1.)#这个操作的目的通常是为了在某些计算中需要使用 self.alphas_cumprod_prev 作为一个与 self.alphas_cumprod 相关的中间变量。在这种情况下添加一个1作为起始值可以确保计算的正确性。# calculations for diffusion q(x_t | x_{t-1}) and others#计算一些用于不同公式的其他变量self.sqrt_alphas_cumprod torch.sqrt(self.alphas_cumprod)self.sqrt_one_minus_alphas_cumprod torch.sqrt(1.0 - self.alphas_cumprod)self.log_one_minus_alphas_cumprod torch.log(1.0 - self.alphas_cumprod)self.sqrt_recip_alphas_cumprod torch.sqrt(1.0 / self.alphas_cumprod)self.sqrt_recipm1_alphas_cumprod torch.sqrt(1.0 / self.alphas_cumprod - 1)# calculations for posterior q(x_{t-1} | x_t, x_0)self.posterior_variance (self.betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod))# below: log calculation clipped because the posterior variance is 0 at the beginning# of the diffusion chain#用于存储后验分布的对数方差self.posterior_log_variance_clipped torch.log(self.posterior_variance.clamp(min 1e-20))#后验均值的系数1self.posterior_mean_coef1 (self.betas * torch.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod))#后验均值的系数2self.posterior_mean_coef2 ((1.0 - self.alphas_cumprod_prev)* torch.sqrt(self.alphas)/ (1.0 - self.alphas_cumprod))def _extract(self, a, t, x_shape):# 辅助函数从a中提取与时间步t对应的参数# get the param of given timestep tbatch_size t.shape[0]out a.to(t.device).gather(0, t).float()#将输出的out的形状改为只有batch_size,其余维度都为1out out.reshape(batch_size, *((1,) * (len(x_shape) - 1)))return outdef q_sample(self, x_start, t, noiseNone):# forward diffusion (using the nice property): q(x_t | x_0)if noise is None:noise torch.randn_like(x_start)#获得第t步的参数数据sqrt_alphas_cumprod_t self._extract(self.sqrt_alphas_cumprod, t, x_start.shape)sqrt_one_minus_alphas_cumprod_t self._extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)#然后和随机产生的噪声进行按比例拟合达到加噪的效果return sqrt_alphas_cumprod_t * x_start sqrt_one_minus_alphas_cumprod_t * noisedef q_mean_variance(self, x_start, t):# Get the mean and variance of q(x_t | x_0).#x_start为需要输入的图像mean self._extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_startvariance self._extract(1.0 - self.alphas_cumprod, t, x_start.shape)log_variance self._extract(self.log_one_minus_alphas_cumprod, t, x_start.shape)return mean, variance, log_variancedef q_posterior_mean_variance(self, x_start, x_t, t):# Compute the mean and variance of the diffusion posterior: q(x_{t-1} | x_t, x_0)posterior_mean (self._extract(self.posterior_mean_coef1, t, x_t.shape) * x_start self._extract(self.posterior_mean_coef2, t, x_t.shape) * x_t)posterior_variance self._extract(self.posterior_variance, t, x_t.shape)posterior_log_variance_clipped self._extract(self.posterior_log_variance_clipped, t, x_t.shape)return posterior_mean, posterior_variance, posterior_log_variance_clipped#反向预测对于输入的x_t反向去噪noisedef predict_start_from_noise(self, x_t, t, noise):# compute x_0 from x_t and pred noise: the reverse of q_samplereturn (self._extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -self._extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise)#最终返回预测的均值方差def p_mean_variance(self, model, x_t, t, clip_denoisedTrue):# compute predicted mean and variance of p(x_{t-1} | x_t)# predict noise using model#unet模块学习加入了时间t这里的t为所有值为t的向量信息的x_t,通过参数调整最终变为我们的反向预测噪声pred_noise model(x_t, t)# get the predicted x_0: different from the algorithm2 in the paper#从反向预测噪声和x_t预测我们的开始值去噪x_recon self.predict_start_from_noise(x_t, t, pred_noise)#将 x_recon 张量中的元素限制在 -1.0 到 1.0 的范围内任何小于 -1.0 的元素都被设置为 -1.0任何大于 1.0 的元素都被设置为 1.0。if clip_denoised:x_recon torch.clamp(x_recon, min-1., max1.)model_mean, posterior_variance, posterior_log_variance self.q_posterior_mean_variance(x_recon, x_t, t)return model_mean, posterior_variance, posterior_log_variancetorch.no_grad()#从最后一步的随机噪声向前进行去噪采样def p_sample(self, model, x_t, t, clip_denoisedTrue):# denoise_step: sample x_{t-1} from x_t and pred_noise# predict mean and variancemodel_mean, _, model_log_variance self.p_mean_variance(model, x_t, t, clip_denoisedclip_denoised)noise torch.randn_like(x_t)# no noise when t 0nonzero_mask ((t ! 0).float().view(-1, *([1] * (len(x_t.shape) - 1))))#判断t是否为0是0则为0非0则为1# compute x_{t-1}pred_img model_mean nonzero_mask * (0.5 * model_log_variance).exp() * noisereturn pred_imgtorch.no_grad()def p_sample_loop(self, model, shape):# denoise: reverse diffusionbatch_size shape[0]device next(model.parameters()).device# start from pure noise (for each example in the batch)img torch.randn(shape, devicedevice)imgs []#tqdm是python中的一个库用于创建进度条以可视化地显示循环的进度。它可以帮助你了解循环还需要多长时间完成特别是在处理大数据集或长时间运行的任务时非常有用。total是定义总的步数#采样传入的image为随机生成的噪声也就代表了最后的x_t时的噪声for i in tqdm(reversed(range(0, timesteps)), descsampling loop time step, totaltimesteps):#torch.full((batch_size,), i)创建一个值都为i的向量img self.p_sample(model, img, torch.full((batch_size,), i, devicedevice, dtypetorch.long))imgs.append(img.cpu().numpy())return imgstorch.no_grad()def sample(self, model, image_size, batch_size8, channels3):# sample new imagesreturn self.p_sample_loop(model, shape(batch_size, channels, image_size, image_size))def train_losses(self, model, x_start, t):# compute train losses# generate random noise# 随机生成一个正态分布noise torch.randn_like(x_start)# get x_t#输入的图像作为x_start正太分布噪声采用我们自己随机生成的#通过前向加噪对输入图像加入t时刻的噪声前向加入噪的噪声作为基准噪声x_noisy self.q_sample(x_start, t, noisenoise)#通过unet对前向生成的噪声和t生成我们的预测噪声predicted_noise model(x_noisy, t)#损失函数就是生成的噪声和预测的噪声进行损失的计算loss F.mse_loss(noise, predicted_noise)return loss看看效果
from PIL import Image
import requests
import matplotlib.pyplot as plt
from torchvision import datasets, transforms%matplotlib inlineurl http://images.cocodataset.org/val2017/000000039769.jpg
image Image.open(requests.get(url, streamTrue).raw)
# image Image.open(/data/000000039769.jpg)image_size 128
transform transforms.Compose([transforms.Resize(image_size),transforms.CenterCrop(image_size),transforms.PILToTensor(),transforms.ConvertImageDtype(torch.float),transforms.Normalize(mean[0.5, 0.5, 0.5], std[0.5, 0.5, 0.5]),
])x_start transform(image).unsqueeze(0)gaussian_diffusion GaussianDiffusion(timesteps500)plt.figure(figsize(16, 8))
for idx, t in enumerate([0, 50, 100, 200, 499]):#根据x_start和t生成从0~t加噪后的结果x_noisy gaussian_diffusion.q_sample(x_start, ttorch.tensor([t]))#squeeze(): 这是一个挤压操作它用于去除输入张量 中维度为1的维度以简化张量的形状#permute(1, 2, 0): 这是一个维度置换操作将第一个维度移到最后一个维度#最后对每个张量1然后乘以127.5原来的数为-1~11变为0~2x127.5变为0~255noisy_image (x_noisy.squeeze().permute(1, 2, 0) 1) * 127.5noisy_image noisy_image.numpy().astype(np.uint8)plt.subplot(1, 5, 1 idx)plt.imshow(noisy_image)plt.axis(off)plt.title(ft{t})②训练
准备数据集
batch_size 64
timesteps 500transform transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean[0.5], std[0.5])
])# use MNIST dataset
dataset datasets.MNIST(./data, trainTrue, downloadTrue, transformtransform)
train_loader torch.utils.data.DataLoader(dataset, batch_sizebatch_size, shuffleTrue)模型
# define model and diffusion
device cuda if torch.cuda.is_available() else cpu
#这里初始化unet模块输入输出的channel为1.注意力模块这里没有加
model UNetModel(in_channels1,model_channels96,out_channels1,channel_mult(1, 2, 2),attention_resolutions[]
)
model.to(device)
#初始化高斯扩散模型只初始化了需要迭代的步骤为500步时间步默认为线性生成的时间步
gaussian_diffusion GaussianDiffusion(timestepstimesteps)
#优化器对unet模型的参数进行优化
optimizer torch.optim.Adam(model.parameters(), lr5e-4)开始训练
epochs 10
for epoch in range(epochs):for step, (images, labels) in enumerate(train_loader):optimizer.zero_grad()batch_size images.shape[0]images images.to(device)# sample t uniformally for every example in the batch#随机生成batch_size个0~timesteps的t对于每次训练数据我们是随机对第其中一个t时刻的加噪过程进行训练和预测t torch.randint(0, timesteps, (batch_size,), devicedevice).long()#输入unet模型样本图像和t计算损失loss gaussian_diffusion.train_losses(model, images, t)#先随机生成一个正太分布(作为我们的加噪的正太分布#将输入的图像images作为x_start#通过前向加噪对输入的图像加入t时刻的噪声此时生成的噪声作为我们的基准噪声#通过unet输入上一步的基准噪声和时间步t我们进行对基准噪声的预测#损失函数计算的就是我们的预测噪声和基准噪声之间的差距采用的是每个像素点的均方差的计算if step % 200 0:print(Loss:, loss.item())#每次训练模型都是让我们的unet模型的参数进行优化让我们的unet模型最终可以根据给定一个加噪了t次后的图像和t去生成一个对于这个基准噪声的预测。也就是我们的unet模型能生成和加入的噪声十分相似的噪声loss.backward()optimizer.step()Loss: 1.2879185676574707
Loss: 0.05010918155312538
Loss: 0.037472739815711975
Loss: 0.03259456530213356
Loss: 0.03238191455602646
Loss: 0.03526081144809723
Loss: 0.019976193085312843
Loss: 0.026588361710309982
Loss: 0.02474384568631649
Loss: 0.025454936549067497
Loss: 0.01776018552482128
Loss: 0.028406977653503418
Loss: 0.026149388402700424
Loss: 0.023932695388793945
Loss: 0.0222737155854702
Loss: 0.025710856541991234
Loss: 0.026215054094791412
Loss: 0.02046349085867405
Loss: 0.02683963067829609
Loss: 0.023800114169716835
Loss: 0.024538405239582062
Loss: 0.021686285734176636
Loss: 0.019745750352740288
Loss: 0.02584003284573555
Loss: 0.026672476902604103
Loss: 0.023941144347190857
Loss: 0.03131483495235443
Loss: 0.018094774335622787
Loss: 0.025758417323231697
Loss: 0.025309113785624504
Loss: 0.0224548801779747
Loss: 0.021184200420975685
Loss: 0.01910235919058323
Loss: 0.024598510935902596
Loss: 0.024002162739634514
Loss: 0.0232978705316782
Loss: 0.016557812690734863
Loss: 0.019946767017245293
Loss: 0.020528556779026985
Loss: 0.01813691109418869
Loss: 0.020777976140379906
Loss: 0.021010225638747215
Loss: 0.02573891542851925
Loss: 0.02588081546127796
Loss: 0.016215061768889427
Loss: 0.025008078664541245
Loss: 0.01972994953393936
Loss: 0.021410418674349785
Loss: 0.024027982726693153
Loss: 0.021927889436483383③测试
generated_images gaussian_diffusion.sample(model, 28, batch_size64, channels1)
# generated_images: [timesteps, batch_size64, channels1, height28, width28]sampling loop time step: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 500/500 [00:3000:00, 16.61it/s]# generate new images
fig plt.figure(figsize(12, 12), constrained_layoutTrue)
#并定义了一个网格布局该布局包含 8 行和 8 列的子图
gs fig.add_gridspec(8, 8)#[-1]表示生成图像的最后一个也就是x0最后生成的图片将数组重新排列为8,828,28的形式
imgs generated_images[-1].reshape(8, 8, 28, 28)
for n_row in range(8):for n_col in range(8):#将图像加入8*8网格对应的位置f_ax fig.add_subplot(gs[n_row, n_col])#将图像的值变换到0~255进行可视化f_ax.imshow((imgs[n_row, n_col]1.0) * 255 / 2, cmapgray)f_ax.axis(off)可以看到我们的扩散模型生成的图像与minist数据集还是非常相似的
展示降噪的过程
# show the denoise steps
fig plt.figure(figsize(12, 12), constrained_layoutTrue)
gs fig.add_gridspec(16, 16)#也就是我们生成的generated_images是一个多维的矩阵stepbatchsize28,28,1 ; 然后我们需要对第i个step过程取其中的第n_row个图片然后去展示这个去噪的过程
for n_row in range(16):for n_col in range(16):f_ax fig.add_subplot(gs[n_row, n_col])#t_idx计算为第几步的噪声从500开始到0t_idx (timesteps // 16) * n_col if n_col 15 else -1#n_now为第n个图像img generated_images[t_idx][n_row].reshape(28, 28)f_ax.imshow((img1.0) * 255 / 2, cmapgray)f_ax.axis(off)三主要训练过程的解析 不好理解之处对于最后的sample阶段的取样过程是如何从随机的噪声一步一步去噪恢复原图像的 流程上讲 如上图所示每次第t时刻我们首先将t时刻的噪声xt和t时刻位置的正弦编码输入unet网络得到我们预测的噪声然后经过对预测的噪声进行处理得到我们预测的均值和方差然后通过参数重参化技巧构建我们生成的预测去噪图像这里我们每次得到的预测结果是作为下一次新的t-1时刻的噪声xt-1然后通过连续的迭代最终生成初始x0时刻的图像也就是我们最终的反向去噪图像。 难点在于 对于均值和方差的预测以及预测图像的重构等过程的数学推导可以先有个大概的印象后续在慢慢攻克
