Mindspore框架CycleGAN模型实现图像风格迁移|（四）CycleGAN模型训练

Mindspore框架：CycleGAN模型实现图像风格迁移算法

1. Mindspore框架CycleGAN模型实现图像风格迁移|（一）CycleGAN神经网络模型构建
2. Mindspore框架CycleGAN模型实现图像风格迁移|（二）实例数据集（苹果2橘子）
3. Mindspore框架CycleGAN模型实现图像风格迁移|（三）损失函数计算
4. Mindspore框架CycleGAN模型实现图像风格迁移|（四）CycleGAN模型训练
5. 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!')