昇思MindSpore进阶教程--Diffusion扩散模型(下)

大家好，我是刘明，明志科技创始人，华为昇思MindSpore布道师。
技术上主攻前端开发、鸿蒙开发和AI算法研究。
努力为大家带来持续的技术分享，如果你也喜欢我的文章，就点个关注吧

数据准备与处理

在这里我们定义一个正则数据集。数据集可以来自简单的真实数据集的图像组成，如Fashion-MNIST、CIFAR-10或ImageNet，其中线性缩放为
。

每个图像的大小都会调整为相同的大小。有趣的是，图像也是随机水平翻转的。根据论文内容：我们在CIFAR10的训练中使用了随机水平翻转；我们尝试了有翻转和没有翻转的训练，并发现翻转可以稍微提高样本质量。

本实验我们选用Fashion_MNIST数据集，我们使用download下载并解压Fashion_MNIST数据集到指定路径。此数据集由已经具有相同分辨率的图像组成，即28x28。

# 下载MNIST数据集
url = 'https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/dataset.zip'
path = download(url, './', kind="zip", replace=True)

from mindspore.dataset import FashionMnistDataset

image_size = 28
channels = 1
batch_size = 16

fashion_mnist_dataset_dir = "./dataset"
dataset = FashionMnistDataset(dataset_dir=fashion_mnist_dataset_dir, usage="train", num_parallel_workers=cpu_count(), shuffle=True, num_shards=1, shard_id=0)

接下来，我们定义一个transform操作，将在整个数据集上动态应用该操作。该操作应用一些基本的图像预处理：随机水平翻转、重新调整，最后使它们的值在[-1,1]范围内。

transforms = [
    RandomHorizontalFlip(),
    ToTensor(),
    lambda t: (t * 2) - 1
]


dataset = dataset.project('image')
dataset = dataset.shuffle(64)
dataset = dataset.map(transforms, 'image')
dataset = dataset.batch(16, drop_remainder=True)

采样

由于我们将在训练期间从模型中采样（以便跟踪进度），我们定义了下面的代码。采样在本文中总结为算法2：

从扩散模型生成新图像是通过反转扩散过程来实现的：我们从
开始，我们从高斯分布中采样纯噪声，然后使用我们的神经网络逐渐去噪（使用它所学习的条件概率），直到我们最终在时间步
结束。如上图所示，我们可以通过使用我们的噪声预测器插入平均值的重新参数化，导出一个降噪程度较低的图像
。请注意，方差是提前知道的。

理想情况下，我们最终会得到一个看起来像是来自真实数据分布的图像。

下面的代码实现了这一点。

def p_sample(model, x, t, t_index):
    betas_t = extract(betas, t, x.shape)
    sqrt_one_minus_alphas_cumprod_t = extract(
        sqrt_one_minus_alphas_cumprod, t, x.shape
    )
    sqrt_recip_alphas_t = extract(sqrt_recip_alphas, t, x.shape)
    model_mean = sqrt_recip_alphas_t * (x - betas_t * model(x, t) / sqrt_one_minus_alphas_cumprod_t)

    if t_index == 0:
        return model_mean
    posterior_variance_t = extract(posterior_variance, t, x.shape)
    noise = randn_like(x)
    return model_mean + ops.sqrt(posterior_variance_t) * noise

def p_sample_loop(model, shape):
    b = shape[0]
    # 从纯噪声开始
    img = randn(shape, dtype=None)
    imgs = []

    for i in tqdm(reversed(range(0, timesteps)), desc='sampling loop time step', total=timesteps):
        img = p_sample(model, img, ms.numpy.full((b,), i, dtype=mstype.int32), i)
        imgs.append(img.asnumpy())
    return imgs

def sample(model, image_size, batch_size=16, channels=3):
    return p_sample_loop(model, shape=(batch_size, channels, image_size, image_size))

请注意，上面的代码是原始实现的简化版本。

训练过程

下面，我们开始训练吧！

# 定义动态学习率
lr = nn.cosine_decay_lr(min_lr=1e-7, max_lr=1e-4, total_step=10*3750, step_per_epoch=3750, decay_epoch=10)

# 定义 Unet模型
unet_model = Unet(
    dim=image_size,
    channels=channels,
    dim_mults=(1, 2, 4,)
)

name_list = []
for (name, par) in list(unet_model.parameters_and_names()):
    name_list.append(name)
i = 0
for item in list(unet_model.trainable_params()):
    item.name = name_list[i]
    i += 1

# 定义优化器
optimizer = nn.Adam(unet_model.trainable_params(), learning_rate=lr)
loss_scaler = DynamicLossScaler(65536, 2, 1000)

# 定义前向过程
def forward_fn(data, t, noise=None):
    loss = p_losses(unet_model, data, t, noise)
    return loss

# 计算梯度
grad_fn = ms.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=False)

# 梯度更新
def train_step(data, t, noise):
    loss, grads = grad_fn(data, t, noise)
    optimizer(grads)
    return loss

import time

epochs = 10

iterator = dataset.create_tuple_iterator(num_epochs=epochs)
for epoch in range(epochs):
    begin_time = time.time()
    for step, batch in enumerate(iterator):
        unet_model.set_train()
        batch_size = batch[0].shape[0]
        t = randint(0, timesteps, (batch_size,), dtype=ms.int32)
        noise = randn_like(batch[0])
        loss = train_step(batch[0], t, noise)

        if step % 500 == 0:
            print(" epoch: ", epoch, " step: ", step, " Loss: ", loss)
    end_time = time.time()
    times = end_time - begin_time
    print("training time:", times, "s")
    # 展示随机采样效果
    unet_model.set_train(False)
    samples = sample(unet_model, image_size=image_size, batch_size=64, channels=channels)
    plt.imshow(samples[-1][5].reshape(image_size, image_size, channels), cmap="gray")
print("Training Success!")

推理过程（从模型中采样）

要从模型中采样，我们可以只使用上面定义的采样函数：

# 采样64个图片
unet_model.set_train(False)
samples = sample(unet_model, image_size=image_size, batch_size=64, channels=channels)

# 展示一个随机效果
random_index = 5
plt.imshow(samples[-1][random_index].reshape(image_size, image_size, channels), cmap="gray")

可以看到这个模型能产生一件衣服！

请注意，我们训练的数据集分辨率相当低（28x28）。

我们还可以创建去噪过程的gif：

import matplotlib.animation as animation

random_index = 53

fig = plt.figure()
ims = []
for i in range(timesteps):
    im = plt.imshow(samples[i][random_index].reshape(image_size, image_size, channels), cmap="gray", animated=True)
    ims.append([im])

animate = animation.ArtistAnimation(fig, ims, interval=50, blit=True, repeat_delay=100)
animate.save('diffusion.gif')
plt.show()

总结

请注意，DDPM论文表明扩散模型是（非）条件图像有希望生成的方向。自那以后，diffusion得到了（极大的）改进，最明显的是文本条件图像生成。下面，我们列出了一些重要的（但远非详尽无遗的）后续工作：

• 改进的去噪扩散概率模型(Nichol et al., 2021)：发现学习条件分布的方差（除平均值外）有助于提高性能

• 用于高保真图像生成的级联扩散模型([Ho et al., 2021)：引入级联扩散，它包括多个扩散模型的流水线，这些模型生成分辨率提高的图像，用于高保真图像合成

• 扩散模型在图像合成上击败了GANs(Dhariwal et al., 2021)：表明扩散模型通过改进U-Net体系结构以及引入分类器指导，可以获得优于当前最先进的生成模型的图像样本质量

• 无分类器扩散指南([Ho et al., 2021)：表明通过使用单个神经网络联合训练条件和无条件扩散模型，不需要分类器来指导扩散模型

• 具有CLIP Latents (DALL-E 2) 的分层文本条件图像生成 (Ramesh et al., 2022)：在将文本标题转换为CLIP图像嵌入之前使用，然后扩散模型将其解码为图像

• 具有深度语言理解的真实文本到图像扩散模型（ImageGen）(Saharia et al., 2022)：表明将大型预训练语言模型（例如T5）与级联扩散结合起来，对于文本到图像的合成很有效

目前，扩散模型的主要（也许唯一）缺点是它们需要多次正向传递来生成图像（对于像GAN这样的生成模型来说，情况并非如此）。然而，有正在进行中的研究表明只需要10个去噪步骤就能实现高保真生成。