Mindspore框架利用扩散模型DDPM天生高分辨率图像|（三）模型训练与推理实践 ...

利用扩散模型DDPM天生高分辨率图像（天生高保真图像项目实践）

Mindspore框架利用扩散模型DDPM天生高分辨率图像|（一）关于denoising diffusion probabilistic model （DDPM）模型
Mindspore框架利用扩散模型DDPM天生高分辨率图像|（二）数据集准备与处理
Mindspore框架利用扩散模型DDPM天生高分辨率图像|（三）模型训练与推理实践

---

一、扩散模型DDPM模型训练

扩散模型天生新图像是通过反转扩散过程来实现。理想情况下，我们最终会得到一个看起来像是来自真实数据分布的图像。

1. def p_sample(model, x, t, t_index):
2.     betas_t = extract(betas, t, x.shape)
3.     sqrt_one_minus_alphas_cumprod_t = extract(
4.         sqrt_one_minus_alphas_cumprod, t, x.shape
5.     )
6.     sqrt_recip_alphas_t = extract(sqrt_recip_alphas, t, x.shape)
7.     model_mean = sqrt_recip_alphas_t * (x - betas_t * model(x, t) / sqrt_one_minus_alphas_cumprod_t)
9.     if t_index == 0:
10.         return model_mean
11.     posterior_variance_t = extract(posterior_variance, t, x.shape)
12.     noise = randn_like(x)
13.     return model_mean + ops.sqrt(posterior_variance_t) * noise
15. def p_sample_loop(model, shape):
16.     b = shape[0]
17.     # 从纯噪声开始
18.     img = randn(shape, dtype=None)
19.     imgs = []
21.     for i in tqdm(reversed(range(0, timesteps)), desc='sampling loop time step', total=timesteps):
22.         img = p_sample(model, img, ms.numpy.full((b,), i, dtype=mstype.int32), i)
23.         imgs.append(img.asnumpy())
24.     return imgs
26. def sample(model, image_size, batch_size=16, channels=3):
27.     return p_sample_loop(model, shape=(batch_size, channels, image_size, image_size))

复制代码

1.1 模型初始化

1. # 定义动态学习率
2. lr = nn.cosine_decay_lr(min_lr=1e-7, max_lr=1e-4, total_step=10*3750, step_per_epoch=3750, decay_epoch=10)
4. # 定义 Unet模型
5. unet_model = Unet(
6.     dim=image_size,
7.     channels=channels,
8.     dim_mults=(1, 2, 4,)
9. )
11. name_list = []
12. for (name, par) in list(unet_model.parameters_and_names()):
13.     name_list.append(name)
14. i = 0
15. for item in list(unet_model.trainable_params()):
16.     item.name = name_list[i]
17.     i += 1
19. # 定义优化器
20. optimizer = nn.Adam(unet_model.trainable_params(), learning_rate=lr)
21. loss_scaler = DynamicLossScaler(65536, 2, 1000)
23. # 定义前向过程
24. def forward_fn(data, t, noise=None):
25.     loss = p_losses(unet_model, data, t, noise)
26.     return loss
28. # 计算梯度
29. grad_fn = ms.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=False)
31. # 梯度更新
32. def train_step(data, t, noise):
33.     loss, grads = grad_fn(data, t, noise)
34.     optimizer(grads)
35.     return loss

复制代码

1.2 模型训练

1. import time
3. epochs = 10
5. iterator = dataset.create_tuple_iterator(num_epochs=epochs)
6. for epoch in range(epochs):
7.     begin_time = time.time()
8.     for step, batch in enumerate(iterator):
9.         unet_model.set_train()
10.         batch_size = batch[0].shape[0]
11.         t = randint(0, timesteps, (batch_size,), dtype=ms.int32)
12.         noise = randn_like(batch[0])
13.         loss = train_step(batch[0], t, noise)
15.         if step % 500 == 0:
16.             print(" epoch: ", epoch, " step: ", step, " Loss: ", loss)
17.     end_time = time.time()
18.     times = end_time - begin_time
19.     print("training time:", times, "s")
20.     # 展示随机采样效果
21.     unet_model.set_train(False)
22.     samples = sample(unet_model, image_size=image_size, batch_size=64, channels=channels)
23.     plt.imshow(samples[-1][5].reshape(image_size, image_size, channels), cmap="gray")
24. print("Training Success!")

复制代码

二、模型推理预测

2.1 推理过程

从模型中采样

1. # 采样64个图片
2. unet_model.set_train(False)
3. samples = sample(unet_model, image_size=image_size, batch_size=64, channels=channels)
4. # 展示一个随机效果
5. # random_index = 5
6. # plt.imshow(samples[-1][random_index].reshape(image_size, image_size, channels), cmap="gray")

复制代码

创建去噪过程

1. import matplotlib.animation as animation
3. random_index = 53
5. fig = plt.figure()
6. ims = []
7. for i in range(timesteps):
8.     im = plt.imshow(samples[i][random_index].reshape(image_size, image_size, channels), cmap="gray", animated=True)
9.     ims.append([im])
11. animate = animation.ArtistAnimation(fig, ims, interval=50, blit=True, repeat_delay=100)
12. animate.save('diffusion.gif')
13. plt.show()

复制代码

三、参考文献

• The Annotated Diffusion Model
  
• 由浅入深了解Diffusion Model
  
• Diffusion扩散模型
  

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。