8行代码教你创建自己第一个扩散模型！

1. Pytorch中的扩散模型

虽然扩散模型还没有普及到与机器学习中的其他旧架构/方法相同的程度，但仍然有可供使用的实现。在PyTorch中使用扩散模型最简单的方法是使用denoising-diffusion-pytorch包，它实现了一个像本文中讨论的图像扩散模型。要安装该软件包，只需在终端中键入以下命令:

pip install denoising_diffusion_pytorch

2. 最小示例

为了训练模型并生成图像，我们首先导入必要的包:

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模型以及几个参数——要生成的图像的大小，扩散过程中的时间步长数，以及在L1和L2规范之间的选择。推荐课程：面向三维视觉的Linux嵌入式系统教程[理论+代码+实战]。

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()

一旦模型训练完毕，我们就可以通过使用sample()的方法diffusion对象。这里，我们生成了4幅图像，这些图像只是噪声，因为我们的训练数据是随机的:

sampled_images = diffusion.sample(batch_size = 4)

3. 训练自定义数据

这denoising-diffusion-pytorch软件包还允许您在特定数据集上训练扩散模型。只需替换'path/to/your/images'中数据集目录路径的字符串Trainer()对象，然后更改image_size设置为适当的值。之后，只需运行代码来训练模型，然后像以前一样进行采样。注意，PyTorch必须在启用CUDA的情况下编译，以便使用Trainer类别:

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数字的渐进去噪，类似于反向扩散:

4. 最后

扩散模型是一种概念上简单而优雅的数据生成方法。他们最先进的成果与非对抗性训练相结合，使他们达到了很高的水平，鉴于他们的新生状态，预计在未来几年还会有进一步的改进。特别是，已经发现扩散模型对于尖端模型的性能至关重要。

本文仅做学术分享，如有侵权，请联系删文。