Mindspore框架:CycleGAN模型实现图像风格迁移算法
- Mindspore框架CycleGAN模型实现图像风格迁移|(一)CycleGAN神经网络模型构建
- Mindspore框架CycleGAN模型实现图像风格迁移|(二)实例数据集(苹果2橘子)
- Mindspore框架CycleGAN模型实现图像风格迁移|(三)损失函数计算
- Mindspore框架CycleGAN模型实现图像风格迁移|(四)CycleGAN模型训练
- Mindspore框架CycleGAN模型实现图像风格迁移|(五)CycleGAN模型推理与资源下载
CycleGAN模型训练
1. CycleGAN模型
流程:将苹果x输入G,得假橘子G(x);将假橘子G(x)和真橘子y输入判别器Dx(Dx为橘子判别器)判断真伪。
将真橘子y输入F,得假苹果F(x);将假苹果F(x)和真苹果x输入判别器Dy(Dy为苹果判别器)判断真伪。
然后在上述过程中,有4个反向传播,分别为G(x)、Dx、F(y)、Dy四个网络的反向传播。
2.构建模型:生成器模型和判别器模型构建
Mindspore框架CycleGAN模型实现图像风格迁移|(一)CycleGAN神经网络模型构建
3.加载在数据集
Mindspore框架CycleGAN模型实现图像风格迁移|(二)实例数据集(苹果2橘子)
4.创建损失函数和优化器
表示真假苹果损失,真假橘子损失,周期一致损失的综合损失函数。
Mindspore框架CycleGAN模型实现图像风格迁移|(三)损失函数计算,优化器选择,模型前向计算损失的过程,计算梯度和反向传播
5.训练模型
import os
import time
import random
import numpy as np
from PIL import Image
from mindspore import Tensor, save_checkpoint
from mindspore import dtype
# 由于时间原因,epochs设置为1,可根据需求进行调整
epochs = 1
save_step_num = 80
save_checkpoint_epochs = 1
save_ckpt_dir = './train_ckpt_outputs/'
print('Start training!')
for epoch in range(epochs):
g_loss = []
d_loss = []
start_time_e = time.time()
for step, data in enumerate(dataset.create_dict_iterator()):
start_time_s = time.time()
img_a = data["image_A"]
img_b = data["image_B"]
res_g = train_step_g(img_a, img_b)
fake_a = res_g[0]
fake_b = res_g[1]
res_d = train_step_d(img_a, img_b, image_pool(fake_a), image_pool(fake_b))
loss_d = float(res_d.asnumpy())
step_time = time.time() - start_time_s
res = []
for item in res_g[2:]:
res.append(float(item.asnumpy()))
g_loss.append(res[0])
d_loss.append(loss_d)
if step % save_step_num == 0:
print(f"Epoch:[{int(epoch + 1):>3d}/{int(epochs):>3d}], "
f"step:[{int(step):>4d}/{int(datasize):>4d}], "
f"time:{step_time:>3f}s,\n"
f"loss_g:{res[0]:.2f}, loss_d:{loss_d:.2f}, "
f"loss_g_a: {res[1]:.2f}, loss_g_b: {res[2]:.2f}, "
f"loss_c_a: {res[3]:.2f}, loss_c_b: {res[4]:.2f}, "
f"loss_idt_a: {res[5]:.2f}, loss_idt_b: {res[6]:.2f}")
epoch_cost = time.time() - start_time_e
per_step_time = epoch_cost / datasize
mean_loss_d, mean_loss_g = sum(d_loss) / datasize, sum(g_loss) / datasize
print(f"Epoch:[{int(epoch + 1):>3d}/{int(epochs):>3d}], "
f"epoch time:{epoch_cost:.2f}s, per step time:{per_step_time:.2f}, "
f"mean_g_loss:{mean_loss_g:.2f}, mean_d_loss:{mean_loss_d :.2f}")
if epoch % save_checkpoint_epochs == 0:
os.makedirs(save_ckpt_dir, exist_ok=True)
save_checkpoint(net_rg_a, os.path.join(save_ckpt_dir, f"g_a_{epoch}.ckpt"))
save_checkpoint(net_rg_b, os.path.join(save_ckpt_dir, f"g_b_{epoch}.ckpt"))
save_checkpoint(net_d_a, os.path.join(save_ckpt_dir, f"d_a_{epoch}.ckpt"))
save_checkpoint(net_d_b, os.path.join(save_ckpt_dir, f"d_b_{epoch}.ckpt"))
print('End of training!')