昇思25天学习打卡营第41天|生成式-Diffusion扩散模型

昇思25天学习打卡营第41天|生成式-Diffusion扩散模型

Diffusion扩散模型

扩散模型（Diffusion Models）是一类基于概率的生成模型，通过一系列逐步的随机噪声注入和去噪过程生成数据。模型的基本思想是先定义一个从数据分布到高斯噪声分布的前向过程（扩散过程），然后学习一个反向过程（去噪过程）以将噪声恢复为原始数据。通过这种方式，扩散模型可以生成逼真的图像、音频等数据。其优势在于生成质量高且训练过程稳定，相较于传统生成对抗网络（GAN）更容易训练，同时具有较强的理论可解释性和灵活性。

扩散过程（Forward Process）和反向过程（Reverse Process）

扩散过程的主要任务：给定原始数据 x0（例如图像），数据在每一步 t被注入噪声，得到一个新的表示 xt，直到最后一步 T ，数据变成近似于标准高斯噪声的分布。

反向过程的主要任务：如何从噪声数据 xT 逐步去噪回原始数据 x0。这个过程也是一个马尔科夫链，每一步用一个神经网络来近似逆转扩散过程。

#在多个时间步长𝑇（timesteps）中，从实际分布逐渐向图像添加噪声
def linear_beta_schedule(timesteps):
    beta_start = 0.0001
    beta_end = 0.02
    return np.linspace(beta_start, beta_end, timesteps).astype(np.float32)
# 扩散200步
timesteps = 200

# 定义 beta schedule
betas = linear_beta_schedule(timesteps=timesteps)

# 定义 alphas
alphas = 1. - betas
alphas_cumprod = np.cumprod(alphas, axis=0)
alphas_cumprod_prev = np.pad(alphas_cumprod[:-1], (1, 0), constant_values=1)

sqrt_recip_alphas = Tensor(np.sqrt(1. / alphas))
sqrt_alphas_cumprod = Tensor(np.sqrt(alphas_cumprod))
sqrt_one_minus_alphas_cumprod = Tensor(np.sqrt(1. - alphas_cumprod))

# 计算 q(x_{t-1} | x_t, x_0)
posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)

p2_loss_weight = (1 + alphas_cumprod / (1 - alphas_cumprod)) ** -0.
p2_loss_weight = Tensor(p2_loss_weight)

def extract(a, t, x_shape):
    b = t.shape[0]
    out = Tensor(a).gather(t, -1)
    return out.reshape(b, *((1,) * (len(x_shape) - 1)))

函数：linear_beta_schedule：生成一个线性增加的噪声调度（beta schedule），用于控制在每个时间步长中添加的噪声量。

alphas：每个时间步的保留比例（数据保留的部分），计算方式为 1 - betas。

posterior_variance：计算在给定时间步长的情况下，从当前状态恢复到前一个状态的方差，用来描述去噪过程中的不确定性。用于确定去噪的噪声量，较大的方差表示去噪过程中的不确定性较大，可能需要更多的步骤或更复杂的模型来进行准确的去噪。

p2_loss_weight：用于训练时对损失的加权调整，有助于在训练过程中平衡模型对不同时间步数据的学习，从而防止模型在某些时间步上过拟合或欠拟合。

损失函数：

def p_losses(unet_model, x_start, t, noise=None):
    if noise is None:
        noise = randn_like(x_start)
    x_noisy = q_sample(x_start=x_start, t=t, noise=noise)
    predicted_noise = unet_model(x_noisy, t)

    loss = nn.SmoothL1Loss()(noise, predicted_noise)# todo
    loss = loss.reshape(loss.shape[0], -1)
    loss = loss * extract(p2_loss_weight, t, loss.shape)
    return loss.mean()

输入数据：

• unet_model: 训练模型，用于预测噪声。
• x_start: 原始的无噪声数据。
• t: 当前时间步。
• noise: 可选的噪声，默认生成随机噪声。

输出数据：

• 返回一个标量值，表示该批次数据的平均损失。这是通过计算预测的噪声与实际添加的噪声之间的差异（误差）得到的。

U-Net神经网络

神经网络需要在特定时间步长接收带噪声的图像，并返回预测的噪声。预测噪声是与输入图像具有相同大小/分辨率的张量。因此，U-Net神经网络通过编码器（下采样路径）和解码器（上采样路径）来在实现对图像特征进行深层次提取同时，保证输出的预测噪声与输入图像具有相同大小/分辨率的张量。

U-Net 在扩散模型中的作用：

噪声预测：

• U-Net 被训练来预测给定噪声图像的原始噪声部分。通过最小化预测噪声与实际噪声之间的差距，网络学习到去除噪声并恢复原始图像的能力。

逆扩散过程：

• 在扩散模型的逆过程（去噪过程）中，U-Net 被用于将含噪图像逐渐还原到无噪或低噪状态。这一步骤对生成高质量的最终图像至关重要。

训练过程

训练扩散模型的损失函数是最小化反向过程与真实扩散过程之间的KL散度。

1. 正向过程模拟：首先通过扩散过程将数据 x0破坏成一系列 xt
2. 反向过程建模：训练一个神经网络去近似每一步的反向条件分布。
3. 损失计算：计算模型生成分布和真实分布之间的KL散度，并将其作为损失。
4. 梯度下降：使用优化算法（如Adam）来调整模型参数，以最小化损失。

# 定义动态学习率
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设置为1，可根据需求进行调整
epochs = 1

for epoch in range(epochs):
    begin_time = time.time()
    for step, batch in enumerate(dataset.create_tuple_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!")