昇思25天学习打卡营第17天|Pix2Pix实现图像转换案例：从环境配置到训练案例的全流程解析

目录

环境配置

数据展示

创建网络

训练

推理

---

环境配置

---

        %%capture captured_output 通常在 Jupyter Notebook 或与之类似的环境中被运用，用于捕获后续代码块的输出，以便在必要时对其进行处理或者予以忽略。首先会尝试卸载已安装的 mindspore 库，接着通过指定的镜像源（https://pypi.mirrors.ustc.edu.cn/simple）来安装指定版本（2.3.0rc1）的 mindspore 库。其还用于查看当前所安装的 mindspore 库的详尽信息，涵盖了版本号等内容。从名为 download 的模块中导入 download 函数。在此部分中，定义了一个 URL 字符串，而后调用导入的 download 函数，从指定的 URL 下载文件并将其保存至./dataset 目录当中。kind="tar" 意味着下载的文件属于 tar 压缩文件类型，replace=True 则表示如果在目标位置存在同名文件，就进行替换操作。

        代码如下：

%%capture captured_output
# 实验环境已经预装了mindspore==2.3.0rc1，如需更换mindspore版本，可更改下面mindspore的版本号
!pip uninstall mindspore -y
!pip install -i https://pypi.mirrors.ustc.edu.cn/simple mindspore==2.3.0rc1
# 查看当前 mindspore 版本
!pip show mindspore
from download import download
url = "https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/models/application/dataset_pix2pix.tar"
download(url, "./dataset", kind="tar", replace=True)

        分析：这段代码主要进行了 mindspore 库的版本管理和指定数据集的下载操作。

数据展示

---

        分别从 mindspore 库导入 dataset 模块并将其别名为 ds ，同时导入 matplotlib.pyplot 模块并别名为 plt ，以用于后续的数据处理及图像可视化操作。创建了一个 MindDataset 对象 dataset ，明确数据文件的路径为 "./dataset/dataset_pix2pix/train.mindrecord" ，并指定需读取的列名为 ["input_images", "target_images"] ，还设置了数据打乱（shuffle=True）。通过 dataset 的 create_dict_iterator 方法创建一个迭代器，再运用 next 函数获取下一个迭代结果，将其存于 data_iter 中，并且设置输出为 numpy 数组的格式。创建一个新的图形窗口，将其大小设定为 (10, 3) ，分辨率设为 140dpi 。借助 for 循环来遍历 data_iter 中 input_images 列的前 10 个图像数据。通过 plt.subplot 划分出子图的位置，关闭坐标轴（plt.axis("off")），接着使用 plt.imshow 来展示图像。在此处对图像进行了特定处理，即 (image.transpose(1, 2, 0) + 1) / 2 ，这或许是为了调整图像的显示效果。最终展示所绘制的图形。

        代码如下：

from mindspore import dataset as ds
import matplotlib.pyplot as plt
dataset = ds.MindDataset("./dataset/dataset_pix2pix/train.mindrecord", columns_list=["input_images", "target_images"], shuffle=True)
data_iter = next(dataset.create_dict_iterator(output_numpy=True))
# 可视化部分训练数据
plt.figure(figsize=(10, 3), dpi=140)
for i, image in enumerate(data_iter['input_images'][:10], 1):
    plt.subplot(3, 10, i)
    plt.axis("off")
    plt.imshow((image.transpose(1, 2, 0) + 1) / 2)
plt.show()

        分析：这段代码主要是读取指定的数据文件，获取其中的图像数据，并进行可视化展示。

        运行结果：

创建网络

---

第一步：定义UNet Skip Connection Block

        导入了 mindspore 库，以及其中的 nn （神经网络模块）和 ops （操作模块）。定义了一个名为 UNetSkipConnectionBlock 的类，此类别继承自 nn.Cell ，意味着这是一个自定义的神经网络单元。def __init__ 这是类的初始化方法，接收多个参数以配置该单元的属性。调用了父类 nn.Cell 的初始化方法。创建了用于下采样和上采样的批归一化层，并设定了 use_bias 的初始值。依据 norm_mode 的值对归一化层的设置以及 use_bias 的值进行调整。处理了 in_planes 参数的默认值。创建了下采样的卷积层、下采样的激活函数 LeakyReLU 以及上采样的激活函数 ReLU 。根据不同的条件（如最外层、最内层或者其他情况）来构建模型的结构，涵盖了下采样部分、上采样部分以及可能存在的 Dropout 层。将构建完成的模型结构存储为 self.model ，并且设置了是否启用跳过连接的标志。def construct 这是前向传播的方法。首先通过 self.model 对输入 x 进行处理，接着依据 self.skip_connections 的值来决定是否执行跳过连接的操作，最终返回输出结果。

        代码如下：

import mindspore
import mindspore.nn as nn
import mindspore.ops as ops

class UNetSkipConnectionBlock(nn.Cell):
    def __init__(self, outer_nc, inner_nc, in_planes=None, dropout=False,
                 submodule=None, outermost=False, innermost=False, alpha=0.2, norm_mode='batch'):
        super(UNetSkipConnectionBlock, self).__init__()
        down_norm = nn.BatchNorm2d(inner_nc)
        up_norm = nn.BatchNorm2d(outer_nc)
        use_bias = False
        if norm_mode == 'instance':
            down_norm = nn.BatchNorm2d(inner_nc, affine=False)
            up_norm = nn.BatchNorm2d(outer_nc, affine=False)
            use_bias = True
        if in_planes is None:
            in_planes = outer_nc
        down_conv = nn.Conv2d(in_planes, inner_nc, kernel_size=4,
                              stride=2, padding=1, has_bias=use_bias, pad_mode='pad')
        down_relu = nn.LeakyReLU(alpha)
        up_relu = nn.ReLU()
        if outermost:
            up_conv = nn.Conv2dTranspose(inner_nc * 2, outer_nc,
                                         kernel_size=4, stride=2,
                                         padding=1, pad_mode='pad')
            down = [down_conv]
            up = [up_relu, up_conv, nn.Tanh()]
            model = down + [submodule] + up
        elif innermost:
            up_conv = nn.Conv2dTranspose(inner_nc, outer_nc,
                                         kernel_size=4, stride=2,
                                         padding=1, has_bias=use_bias, pad_mode='pad')
            down = [down_relu, down_conv]
            up = [up_relu, up_conv, up_norm]
            model = down + up
        else:
            up_conv = nn.Conv2dTranspose(inner_nc * 2, outer_nc,
                                         kernel_size=4, stride=2,
                                         padding=1, has_bias=use_bias, pad_mode='pad')
            down = [down_relu, down_conv, down_norm]
            up = [up_relu, up_conv, up_norm]

            model = down + [submodule] + up
            if dropout:
                model.append(nn.Dropout(p=0.5))
        self.model = nn.SequentialCell(model)
        self.skip_connections = not outermost

    def construct(self, x):
        out = self.model(x)
        if self.skip_connections:
            out = ops.concat((out, x), axis=1)
        return out

        分析：这段代码定义了一个用于 UNet 架构中的具有跳过连接功能的模块，根据不同的条件配置其内部结构，并实现了前向传播的计算逻辑。

第二步：基于UNet的生成器

        定义了一个名为 UNetGenerator 的类，其继承自 nn.Cell 。def __init__ 这是类的初始化方法，接收多个参数以配置生成器的属性。调用了父类 nn.Cell 的初始化方法。创建了最内层的 UNetSkipConnectionBlock 模块。通过循环创建中间层的 UNetSkipConnectionBlock 模块，并将前一个模块作为子模块进行传递。依次创建不同规模的 UNetSkipConnectionBlock 模块，且均将前一个模块作为子模块。创建最外层的 UNetSkipConnectionBlock 模块，并将之前构建的模块作为子模块，最终将其存储为 self.model 。定义了前向传播的方法，直接返回通过 self.model 处理输入 x 的结果。

        代码如下：

class UNetGenerator(nn.Cell):
    def __init__(self, in_planes, out_planes, ngf=64, n_layers=8, norm_mode='bn', dropout=False):
        super(UNetGenerator, self).__init__()
        unet_block = UNetSkipConnectionBlock(ngf * 8, ngf * 8, in_planes=None, submodule=None,
                                             norm_mode=norm_mode, innermost=True)
        for _ in range(n_layers - 5):
            unet_block = UNetSkipConnectionBlock(ngf * 8, ngf * 8, in_planes=None, submodule=unet_block,
                                                 norm_mode=norm_mode, dropout=dropout)
        unet_block = UNetSkipConnectionBlock(ngf * 4, ngf * 8, in_planes=None, submodule=unet_block,
                                             norm_mode=norm_mode)
        unet_block = UNetSkipConnectionBlock(ngf * 2, ngf * 4, in_planes=None, submodule=unet_block,
                                             norm_mode=norm_mode)
        unet_block = UNetSkipConnectionBlock(ngf, ngf * 2, in_planes=None, submodule=unet_block,
                                             norm_mode=norm_mode)
        self.model = UNetSkipConnectionBlock(out_planes, ngf, in_planes=in_planes, submodule=unet_block,
                                             outermost=True, norm_mode=norm_mode)

    def construct(self, x):
        return self.model(x)

        分析：这段代码定义了一个 UNetGenerator 类，通过层层嵌套创建 UNetSkipConnectionBlock 模块来构建生成器的结构，并实现了前向传播的逻辑。

第三步：基于PatchGAN的判别器

        首先，导入了 mindspore.nn 模块并别名为 nn 。

        定义了一个名为 ConvNormRelu 的类，它继承自 nn.Cell 。

        在 ConvNormRelu 类中：

        __init__ 方法用于初始化，根据参数配置卷积层、归一化层和激活函数，并将它们组合成一个序列。

        construct 方法执行前向传播，将输入通过之前构建的序列进行处理并输出。

        然后，定义了一个名为 Discriminator 的类，它也继承自 nn.Cell 。

        在 Discriminator 类中：

        __init__ 方法首先创建了初始的卷积层和激活函数层，然后通过循环添加多个 ConvNormRelu 模块来构建网络结构，最后添加一个输出卷积层。

        construct 方法先将输入 x 和 y 沿通道维度拼接，然后通过构建的网络结构进行处理并输出。

        代码如下：

import mindspore.nn as nn

class ConvNormRelu(nn.Cell):
    def __init__(self,
                 in_planes,
                 out_planes,
                 kernel_size=4,
                 stride=2,
                 alpha=0.2,
                 norm_mode='batch',
                 pad_mode='CONSTANT',
                 use_relu=True,
                 padding=None):
        super(ConvNormRelu, self).__init__()
        norm = nn.BatchNorm2d(out_planes)
        if norm_mode == 'instance':
            norm = nn.BatchNorm2d(out_planes, affine=False)
        has_bias = (norm_mode == 'instance')
        if not padding:
            padding = (kernel_size - 1) // 2
        if pad_mode == 'CONSTANT':
            conv = nn.Conv2d(in_planes, out_planes, kernel_size, stride, pad_mode='pad',
                             has_bias=has_bias, padding=padding)
            layers = [conv, norm]
        else:
            paddings = ((0, 0), (0, 0), (padding, padding), (padding, padding))
            pad = nn.Pad(paddings=paddings, mode=pad_mode)
            conv = nn.Conv2d(in_planes, out_planes, kernel_size, stride, pad_mode='pad', has_bias=has_bias)
            layers = [pad, conv, norm]
        if use_relu:
            relu = nn.ReLU()
            if alpha > 0:
                relu = nn.LeakyReLU(alpha)
            layers.append(relu)
        self.features = nn.SequentialCell(layers)

    def construct(self, x):
        output = self.features(x)
        return output

class Discriminator(nn.Cell):
    def __init__(self, in_planes=3, ndf=64, n_layers=3, alpha=0.2, norm_mode='batch'):
        super(Discriminator, self).__init__()
        kernel_size = 4
        layers = [
            nn.Conv2d(in_planes, ndf, kernel_size, 2, pad_mode='pad', padding=1),
            nn.LeakyReLU(alpha)
        ]
        nf_mult = ndf
        for i in range(1, n_layers):
            nf_mult_prev = nf_mult
            nf_mult = min(2 ** i, 8) * ndf
            layers.append(ConvNormRelu(nf_mult_prev, nf_mult, kernel_size, 2, alpha, norm_mode, padding=1))
        nf_mult_prev = nf_mult
        nf_mult = min(2 ** n_layers, 8) * ndf
        layers.append(ConvNormRelu(nf_mult_prev, nf_mult, kernel_size, 1, alpha, norm_mode, padding=1))
        layers.append(nn.Conv2d(nf_mult, 1, kernel_size, 1, pad_mode='pad', padding=1))
        self.features = nn.SequentialCell(layers)

    def construct(self, x, y):
        x_y = ops.concat((x, y), axis=1)
        output = self.features(x_y)
        return output
分析：这段代码定义了两个用于图像处理的神经网络类，一个是包含卷积、归一化和激活操作的基本模块 ConvNormRelu ，另一个是判别器 Discriminator ，通过组合这些模块来构建网络。

第四步：第四步：Pix2Pix的生成器和判别器初始化

        首先，定义了一些全局的参数，如输入和输出的平面数、生成器和判别器的相关参数、初始化的增益和类型等。

        然后，创建了一个 UNetGenerator 类型的生成器 net_generator ，并对其内部的卷积层和转置卷积层的权重进行初始化，根据 init_type 的不同选择不同的初始化方式，对批归一化层的参数也进行了相应的初始化。

        接着，创建了一个 Discriminator 类型的判别器 net_discriminator ，并进行了与生成器类似的权重初始化操作。

        最后，定义了一个名为 Pix2Pix 的类，它继承自 nn.Cell 。在其初始化方法中，接收判别器和生成器，并将它们作为成员变量。construct 方法用于前向传播，输入 reala ，通过生成器生成 fakeb 并返回。

        代码如下：

import mindspore.nn as nn
from mindspore.common import initializer as init

g_in_planes = 3
g_out_planes = 3
g_ngf = 64
g_layers = 8
d_in_planes = 6
d_ndf = 64
d_layers = 3
alpha = 0.2
init_gain = 0.02
init_type = 'normal'


net_generator = UNetGenerator(in_planes=g_in_planes, out_planes=g_out_planes,
                              ngf=g_ngf, n_layers=g_layers)
for _, cell in net_generator.cells_and_names():
    if isinstance(cell, (nn.Conv2d, nn.Conv2dTranspose)):
        if init_type == 'normal':
            cell.weight.set_data(init.initializer(init.Normal(init_gain), cell.weight.shape))
        elif init_type == 'xavier':
            cell.weight.set_data(init.initializer(init.XavierUniform(init_gain), cell.weight.shape))
        elif init_type == 'constant':
            cell.weight.set_data(init.initializer(0.001, cell.weight.shape))
        else:
            raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
    elif isinstance(cell, nn.BatchNorm2d):
        cell.gamma.set_data(init.initializer('ones', cell.gamma.shape))
        cell.beta.set_data(init.initializer('zeros', cell.beta.shape))


net_discriminator = Discriminator(in_planes=d_in_planes, ndf=d_ndf,
                                  alpha=alpha, n_layers=d_layers)
for _, cell in net_discriminator.cells_and_names():
    if isinstance(cell, (nn.Conv2d, nn.Conv2dTranspose)):
        if init_type == 'normal':
            cell.weight.set_data(init.initializer(init.Normal(init_gain), cell.weight.shape))
        elif init_type == 'xavier':
            cell.weight.set_data(init.initializer(init.XavierUniform(init_gain), cell.weight.shape))
        elif init_type == 'constant':
            cell.weight.set_data(init.initializer(0.001, cell.weight.shape))
        else:
            raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
    elif isinstance(cell, nn.BatchNorm2d):
        cell.gamma.set_data(init.initializer('ones', cell.gamma.shape))
        cell.beta.set_data(init.initializer('zeros', cell.beta.shape))

class Pix2Pix(nn.Cell):
    """Pix2Pix模型网络"""
    def __init__(self, discriminator, generator):
        super(Pix2Pix, self).__init__(auto_prefix=True)
        self.net_discriminator = discriminator
        self.net_generator = generator

    def construct(self, reala):
        fakeb = self.net_generator(reala)
        return fakeb

训练

---

        首先，定义了一些训练相关的参数，如训练的轮数、检查点保存目录、数据集大小、学习率等。

        然后，定义了一个函数 get_lr 用于生成动态变化的学习率。

        加载数据集，并定义了损失函数 loss_f （二分类交叉熵损失）和 l1_loss （L1 损失）。

        定义了两个前向传播函数 forword_dis 用于判别器的损失计算，forword_gan 用于生成器的损失计算。

        创建了判别器和生成器的优化器 d_opt 和 g_opt ，并使用 value_and_grad 计算损失函数关于可训练参数的梯度。

        定义了 train_step 函数，用于执行一步训练，包括计算判别器和生成器的损失，并应用优化器更新参数。

        在训练循环中，遍历数据集进行训练，记录每步的损失，打印训练的进度和损失信息，并在最后一轮训练结束时保存生成器的检查点。

        代码如下：

import numpy as np
import os
import datetime
from mindspore import value_and_grad, Tensor

epoch_num = 3
ckpt_dir = "results/ckpt"
dataset_size = 400
val_pic_size = 256
lr = 0.0002
n_epochs = 100
n_epochs_decay = 100

def get_lr():
    lrs = [lr] * dataset_size * n_epochs
    lr_epoch = 0
    for epoch in range(n_epochs_decay):
        lr_epoch = lr * (n_epochs_decay - epoch) / n_epochs_decay
        lrs += [lr_epoch] * dataset_size
    lrs += [lr_epoch] * dataset_size * (epoch_num - n_epochs_decay - n_epochs)
    return Tensor(np.array(lrs).astype(np.float32))

dataset = ds.MindDataset("./dataset/dataset_pix2pix/train.mindrecord", columns_list=["input_images", "target_images"], shuffle=True, num_parallel_workers=1)
steps_per_epoch = dataset.get_dataset_size()
loss_f = nn.BCEWithLogitsLoss()
l1_loss = nn.L1Loss()

def forword_dis(reala, realb):
    lambda_dis = 0.5
    fakeb = net_generator(reala)
    pred0 = net_discriminator(reala, fakeb)
    pred1 = net_discriminator(reala, realb)
    loss_d = loss_f(pred1, ops.ones_like(pred1)) + loss_f(pred0, ops.zeros_like(pred0))
    loss_dis = loss_d * lambda_dis
    return loss_dis

def forword_gan(reala, realb):
    lambda_gan = 0.5
    lambda_l1 = 100
    fakeb = net_generator(reala)
    pred0 = net_discriminator(reala, fakeb)
    loss_1 = loss_f(pred0, ops.ones_like(pred0))
    loss_2 = l1_loss(fakeb, realb)
    loss_gan = loss_1 * lambda_gan + loss_2 * lambda_l1
    return loss_gan

d_opt = nn.Adam(net_discriminator.trainable_params(), learning_rate=get_lr(),
                beta1=0.5, beta2=0.999, loss_scale=1)
g_opt = nn.Adam(net_generator.trainable_params(), learning_rate=get_lr(),
                beta1=0.5, beta2=0.999, loss_scale=1)

grad_d = value_and_grad(forword_dis, None, net_discriminator.trainable_params())
grad_g = value_and_grad(forword_gan, None, net_generator.trainable_params())

def train_step(reala, realb):
    loss_dis, d_grads = grad_d(reala, realb)
    loss_gan, g_grads = grad_g(reala, realb)
    d_opt(d_grads)
    g_opt(g_grads)
    return loss_dis, loss_gan

if not os.path.isdir(ckpt_dir):
    os.makedirs(ckpt_dir)

g_losses = []
d_losses = []
data_loader = dataset.create_dict_iterator(output_numpy=True, num_epochs=epoch_num)

for epoch in range(epoch_num):
    for i, data in enumerate(data_loader):
        start_time = datetime.datetime.now()
        input_image = Tensor(data["input_images"])
        target_image = Tensor(data["target_images"])
        dis_loss, gen_loss = train_step(input_image, target_image)
        end_time = datetime.datetime.now()
        delta = (end_time - start_time).microseconds
        if i % 2 == 0:
            print("ms per step:{:.2f}  epoch:{}/{}  step:{}/{}  Dloss:{:.4f}  Gloss:{:.4f} ".format((delta / 1000), (epoch + 1), (epoch_num), i, steps_per_epoch, float(dis_loss), float(gen_loss)))
        d_losses.append(dis_loss.asnumpy())
        g_losses.append(gen_loss.asnumpy())
    if (epoch + 1) == epoch_num:
        mindspore.save_checkpoint(net_generator, ckpt_dir + "Generator.ckpt")

        分析：这段代码实现了一个基于给定模型和损失函数的训练过程，包括学习率的调整、损失计算、参数更新和模型保存等操作。

        运行结果：

推理

---

        首先，从指定的检查点文件（ckpt_dir + "Generator.ckpt"）加载生成器的参数，并将其加载到已定义的生成器网络 net_generator 中。

        然后，重新加载数据集。

        通过数据集创建一个迭代器，并获取下一个迭代的数据。

        使用加载了参数的生成器网络对迭代器中的输入图像进行预测，得到预测结果 predict_show 。

        使用 matplotlib.pyplot 绘制图像，展示输入图像和生成器生成的预测图像。前 10 个子图展示输入图像，后 10 个子图展示对应的预测图像。

        最后，显示绘制的图像。

        代码如下：

from mindspore import load_checkpoint, load_param_into_net

param_g = load_checkpoint(ckpt_dir + "Generator.ckpt")
load_param_into_net(net_generator, param_g)
dataset = ds.MindDataset("./dataset/dataset_pix2pix/train.mindrecord", columns_list=["input_images", "target_images"], shuffle=True)
data_iter = next(dataset.create_dict_iterator())
predict_show = net_generator(data_iter["input_images"])
plt.figure(figsize=(10, 3), dpi=140)
for i in range(10):
    plt.subplot(2, 10, i + 1)
    plt.imshow((data_iter["input_images"][i].asnumpy().transpose(1, 2, 0) + 1) / 2)
    plt.axis("off")
    plt.subplots_adjust(wspace=0.05, hspace=0.02)
    plt.subplot(2, 10, i + 11)
    plt.imshow((predict_show[i].asnumpy().transpose(1, 2, 0) + 1) / 2)
    plt.axis("off")
    plt.subplots_adjust(wspace=0.05, hspace=0.02)
plt.show()

        分析：这段代码的主要目的是加载训练好的生成器模型参数，对新的数据进行预测，并可视化输入图像和生成的预测图像。

        运行结果：

       打印时间：