Pytorch中如何使用扩散模型？机器学习扩散模型简介

Posted 2022-12-13

tags:

篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了Pytorch中如何使用扩散模型？机器学习扩散模型简介相关的知识，希望对你有一定的参考价值。

扩散模型是一种生成模型，在过去的几年里忽然火了起来，这其中也是有一定原因的。

单看2020前后的几篇开创性的文章我们就可以知道扩散模型的性能了，比如在图像生成方面打败了GANs。最近，研究人员想必也都看到了OpenAI上个月发布的图像生成模型DALL-E 2中也是用到了扩散模型。

DALL-E 2使用同一个caption生成的不同的图片

鉴于最近扩散模型的热度越来越高，许多机器学习从业者肯定也想知道其内部原理，那今天这个文章，我们将带大家一起研究扩散模型的理论基础，然后演示如何在PyTorch中使用扩散模型生成图像。

让我们开始吧！！!

扩散模型简介

扩散模型是生成模型，这意味着扩散模型可以生成和训练集类似的数据。从根本上讲的话，扩散模型做的是通过连续添加高斯噪声破坏训练数据，然后反向的噪声处理，也就是通过消除噪声，学习如何去重建数据。训练之后，我们只需要为其输入简单的随机采样的噪声，通过扩散模型学习到的去噪过程，就可以生成数据了。

Diffusion Models can be used to generate images from noise

具体来讲，一个扩散模型是一个潜变量模型，是使用固定的马尔科夫链映射到潜在空间。该马尔科夫链逐渐向数据添加噪声，去获得近似的后验分布$q(\\mathbf x_1:T | \\mathbf x_0)$，其中$x_1 \\dots x_T$都是和$x_0$同纬度的潜变量。在下图中我们可以清晰地看到这条马尔科夫链是怎么在图像上工作的。

最终，图像会被近似转化为纯高斯噪声，训练一个扩散模型的目的是学习 反向降噪过程，比如训练$p_\\theta(x_t-1|x_t)$，反向通过这条链路，我们可以获得新的数据。

扩散模型有何优点

和前边说到的一样，近几年扩散模型发展如火如荼。受到非平衡热力学的启发，扩散模型现在取得了SOTA的图像质量，下边我们就可以看几个例子：

除了能生成高质量的图片之外，扩散模型还有一些其他的优点，比如 不需要对抗训练。我说对抗训练难搞可不是空口无凭，所以如果存在非对抗性的训练能取得差不多的效果和训练效率，最好还是选择非对抗性的方法。就训练效率来说，扩散模型还具有可伸缩性和并行性的额外好处。

虽然扩散模型看起来好像结果都是凭空产生的，但其实是有许多谨慎而有趣的数学抉择和细节为这些结果提供了基础。并且在文献中不断涌现出更多好的效果，那接下来我们就看一下支持扩散模型的数学原理。

深入研究扩散模型

前边已经提到过了，一个扩散模型是有两部分组成的：

前向过程，或者称之为扩散过程，在该过程中一个数据被逐渐噪声化。
反向过程，或称为反向扩散过程，在这个过程中噪声从目标分布中逐步转换回样本的样子。

当噪声水平足够低时，正向过程中的采样链转化可以设置为条件高斯。如果将这一特性和马尔可夫假设结合起来，就得到了正向过程的简单参数化：

$$ q\\left(\\mathbfx1: T \\mid \\mathbfx0\\right):=\\prodt=1^T q\\left(\\mathbfxt \\mid \\mathbfxt-1\\right):=\\prodt=1^T \\mathcalN\\left(\\mathbfxt ; \\sqrt1-\\beta_t \\mathbfxt-1, \\beta_t \\mathbfI\\right) $$

其中$\\beta_1 \\dots \\beta_T$是变量表，可以是固定值也可以是可学习参数。如果表现良好，那当$T$足够大的时候，$x_T$是近似于各向同性的高斯分布的。

Given the Markov assumption, the joint distribution of the latent variables is the product of the Gaussian conditional chain transitions

扩散模型的神奇之处在于其反向过程，在训练过程中学习如何逆置湖=扩散过程，然后与生成新的数据。从纯高斯噪声$p(\\mathbfx T):=\\mathcalN(\\mathbfx T, \\mathbf0, \\mathbfI)$开始，模型通过下面的公式去学习联合分布$p \\theta(\\mathbfx 0: T)$：

$$ p_\\theta\\left(\\mathbfx0: T\\right):=p\\left(\\mathbfxT\\right) \\prodt=1^T p\\theta\\left(\\mathbfxt-1 \\mid \\mathbfxt\\right):=p\\left(\\mathbfxT\\right) \\prodt=1^T \\mathcalN\\left(\\mathbfxt-1 ; \\boldsymbol\\mu\\theta\\left(\\mathbfxt, t\\right), \\boldsymbol\\Sigma\\theta\\left(\\mathbfx_t, t\\right)\\right) $$

其中高斯变缓的时间参数是可学参数，注意因为这里反向扩散过程我们认为是符合马尔科夫假设的，所以高斯变缓的分布仅取决于上一步的分布状况（或者说下一步的分布状况，由你是怎么看模型方向决定的）：

$$ p_\\theta\\left(\\mathbfxt-1 \\mid \\mathbfxt\\right):=\\mathcalN\\left(\\mathbfxt-1 ; \\boldsymbol\\mu\\theta\\left(\\mathbfxt, t\\right), \\mathbf\\Sigma\\theta\\left(\\mathbfx_t, t\\right)\\right) $$

扩散模型在PyTorch中如何使用

虽然扩散模型不像其他机器学习中传统的方法和框架一样普及，但是现在也是有一些实现好的扩散模型可以让我们直接使用的。在Pytorch中使用扩散模型最简单的方法就是使用denoising-diffusion-pytorch库，其中实现了一个和我们上边介绍的一样的图像扩散模型。

想用这个库，先要安装嘛，在终端输入：

pip install denoising_diffusion_pytorch

举个小例子

想训练一个模型并生成图像，我们首先要导入必要的库。

import torch
from denoising_diffusion_pytorch import Unet, GaussianDiffusion

然后我们定义好网络结构，在这个例子里我们用的是U-Net，dim参数指的是在第一次下采样之前特征图的数量；dim_mults中存储的是连续下采样的乘数。

model = Unet(
    dim = 64,
    dim_mults = (1, 2, 4, 8)
)

现在我们的U-Net网络结构已经定义好了，我们需要定义扩散模型本身了。我们需要传入以下几个东西：

我们刚才用那几个参数定义好的U-Net模型。
需要产生的图片的尺寸。
扩散模型需要扩散的时间步长。
是选择L1还是L2归一化。

diffusion = GaussianDiffusion(
    model,
    image_size = 128,
    timesteps = 1000,   # number of steps
    loss_type = l1    # L1 or L2
)

现在这样扩散模型就算是定义好了。可以去训练了。我们现在就搞几个随机数来训练一下，演示演示，然后再用常见的方法训练一下扩散模型。

training_images = torch.randn(8, 3, 128, 128)
loss = diffusion(training_images)
loss.backward()

模型一旦训练好之后，我们最后就可以直接使用diffusion.sample()生成图像，下边这段代码就是我们生成四张图。因为我们的训练数据就是随便搞得随机数嘛，所以生成的肯定也只是随机噪声。

sampled_images = diffusion.sample(batch_size = 4)

在自己的数据集上运行

上边我们已经用随机数举例子了，denoising-diffusion-pytorch写都写了，肯定是可以让你用自己的数据集来训练一个扩散模型啊。你只需要简单地替换一下Trainer()中的path/to/your/images字符串，将其改为数据集的文件夹地址，然后将image_size改为合适的图片大小即可。之后你只需要运行代码就可以训练模型了，采样生成图片和上一段讲的都是一样的。

注意， 想要用Trainer你必须确保你运行代码的地方是支持CUDA的。

from denoising_diffusion_pytorch import Unet, GaussianDiffusion, Trainer

model = Unet(
    dim = 64,
    dim_mults = (1, 2, 4, 8)
).cuda()

diffusion = GaussianDiffusion(
    model,
    image_size = 128,
    timesteps = 1000,   # number of steps
    loss_type = l1    # L1 or L2
).cuda()

trainer = Trainer(
    diffusion,
    path/to/your/images,
    train_batch_size = 32,
    train_lr = 2e-5,
    train_num_steps = 700000,         # total training steps
    gradient_accumulate_every = 2,    # gradient accumulation steps
    ema_decay = 0.995,                # exponential moving average decay
    amp = True                        # turn on mixed precision
)

trainer.train()

下边这个动图里你可以看到从多元噪声逐渐去噪最后生成MNIST数据集中图像的一个连续的去噪过程。

结语

扩散模型是一种概念上简单而优雅的方法，可以解决数据生成的问题。扩散模型现在已经可以在效果上取得SOTA，并且使用的是非对抗性训练，因此该模型现在已经到了一个很高的高度。键入扩散模型如今的地位，可以预期在未来几年还会有进一步的提高。尤其是我们可以看到对于DALL-E 2这样的前沿模型来说，扩散模型的性能至关重要。