昇思MindSpore 应用学习-Pix2Pix实现图像转换-CSDN

Pix2Pix实现图像转换 AI代码解析

Pix2Pix概述

Pix2Pix是基于条件生成对抗网络（cGAN, Condition Generative Adversarial Networks ）实现的一种深度学习图像转换模型，该模型是由Phillip Isola等作者在2017年CVPR上提出的，可以实现语义/标签到真实图片、灰度图到彩色图、航空图到地图、白天到黑夜、线稿图到实物图的转换。Pix2Pix是将cGAN应用于有监督的图像到图像翻译的经典之作，其包括两个模型：生成器和判别器。
传统上，尽管此类任务的目标都是相同的从像素预测像素，但每项都是用单独的专用机器来处理的。而Pix2Pix使用的网络作为一个通用框架，使用相同的架构和目标，只在不同的数据上进行训练，即可得到令人满意的结果，鉴于此许多人已经使用此网络发布了他们自己的艺术作品。

基础原理

cGAN的生成器与传统GAN的生成器在原理上有一些区别，cGAN的生成器是将输入图片作为指导信息，由输入图像不断尝试生成用于迷惑判别器的“假”图像，由输入图像转换输出为相应“假”图像的本质是从像素到另一个像素的映射，而传统GAN的生成器是基于一个给定的随机噪声生成图像，输出图像通过其他约束条件控制生成，这是cGAN和GAN的在图像翻译任务中的差异。Pix2Pix中判别器的任务是判断从生成器输出的图像是真实的训练图像还是生成的“假”图像。在生成器与判别器的不断博弈过程中，模型会达到一个平衡点，生成器输出的图像与真实训练数据使得判别器刚好具有50%的概率判断正确。
在教程开始前，首先定义一些在整个过程中需要用到的符号：

• ：代表观测图像的数据。
• ：代表随机噪声的数据。
• ：生成器网络，给出由观测图像与随机噪声生成的“假”图片，其中来自于训练数据而非生成器。
• ：判别器网络，给出图像判定为真实图像的概率，其中来自于训练数据，来自于生成器。

cGAN的目标可以表示为：

该公式是cGAN的损失函数，D想要尽最大努力去正确分类真实图像与“假”图像，也就是使参数最大化；而G则尽最大努力用生成的“假”图像欺骗D，避免被识破，也就是使参数最小化。cGAN的目标可简化为：

为了对比cGAN和GAN的不同，我们将GAN的目标也进行了说明：

从公式可以看出，GAN直接由随机噪声生成“假”图像，不借助观测图像的任何信息。过去的经验告诉我们，GAN与传统损失混合使用是有好处的，判别器的任务不变，依旧是区分真实图像与“假”图像，但是生成器的任务不仅要欺骗判别器，还要在传统损失的基础上接近训练数据。假设cGAN与L1正则化混合使用，那么有:

进而得到最终目标：

图像转换问题本质上其实就是像素到像素的映射问题，Pix2Pix使用完全一样的网络结构和目标函数，仅更换不同的训练数据集就能分别实现以上的任务。本任务将借助MindSpore框架来实现Pix2Pix的应用。

准备环节

配置环境文件

本案例在GPU，CPU和Ascend平台的动静态模式都支持。

from download import download  # 从下载模块导入下载函数

# 定义要下载文件的URL
url = "https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/models/application/dataset_pix2pix.tar"

# 使用download函数下载文件，指定保存路径、文件类型和是否替换已存在的文件
download(url, "./dataset", kind="tar", replace=True)

代码解析：

1. from download import download:
   • 这一行从download模块中导入download函数，用于后续下载文件。
2. url = "https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/models/application/dataset_pix2pix.tar":
   • 这里定义了一个字符串变量url，它包含了要下载的文件的网络地址。
3. download(url, "./dataset", kind="tar", replace=True):
   • 调用download函数开始下载文件。
   • 参数解析：
     • url: 指定要下载的文件的URL地址。
     • "./dataset": 指定下载后文件的保存路径为当前目录下的dataset文件夹。
     • kind="tar": 指定下载文件的类型为tar格式，可能影响文件解压或处理的方式。
     • replace=True: 如果目标路径已经存在同名文件，则允许覆盖该文件。

API解析：

• download函数通常用于方便地从网上下载文件，支持多种文件类型。其参数能够灵活设置下载的目的地、处理方式及是否覆盖已有的文件。

数据展示

调用Pix2PixDataset和create_train_dataset读取训练集，这里我们直接下载已经处理好的数据集。

from mindspore import dataset as ds  # 从mindspore中导入dataset模块
import matplotlib.pyplot as plt  # 导入matplotlib库中的pyplot模块，用于绘图

# 创建MindDataset对象，加载指定路径的MindRecord数据集
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):  # 遍历前10张输入图像
    plt.subplot(3, 10, i)  # 创建3行10列的子图
    plt.axis("off")  # 不显示坐标轴
    plt.imshow((image.transpose(1, 2, 0) + 1) / 2)  # 显示图像，并进行归一化处理
plt.show()  # 展示所有子图

代码解析：

1. from mindspore import dataset as ds:
   • 从MindSpore框架中导入数据集模块，并将其简化为ds，以便后续使用。
2. import matplotlib.pyplot as plt:
   • 导入matplotlib库的pyplot模块，常用于生成各种图表和可视化。
3. dataset = ds.MindDataset("./dataset/dataset_pix2pix/train.mindrecord", columns_list=["input_images", "target_images"], shuffle=True):
   • 创建一个MindDataset对象，使用指定路径的MindRecord文件。
   • columns_list参数指定要加载的列，这里为input_images和target_images，即输入图像和目标图像。
   • shuffle=True表示在加载数据时打乱顺序。
4. data_iter = next(dataset.create_dict_iterator(output_numpy=True)):
   • 创建一个字典迭代器，提取数据集中的一项数据，并将输出格式设为NumPy数组。
5. plt.figure(figsize=(10, 3), dpi=140):
   • 初始化一个新的图形，设定其大小为10x3英寸，分辨率为140 DPI。
6. for i, image in enumerate(data_iter['input_images'][:10], 1)::
   • 遍历提取的前10张输入图像，通过enumerate获取索引和图像数据。
7. plt.subplot(3, 10, i):
   • 在图形中创建一个3行10列的子图布局，i指定当前图像在布局中的位置。
8. plt.axis("off"):
   • 关闭当前子图的坐标轴显示。
9. plt.imshow((image.transpose(1, 2, 0) + 1) / 2):
   • 将图像的维度从(C, H, W)转换为(H, W, C)，并进行归一化处理，使像素值范围保持在[0, 1]之间。
10. plt.show():
    • 显示所有子图。

API解析：

• MindDataset:
  • 是MindSpore提供的用于处理MindRecord格式数据集的类，支持列选择、数据打乱等功能。
• create_dict_iterator:
  • 用于创建一个字典迭代器，便于按需提取数据集中的数据项。
• output_numpy=True:
  • 指定输出格式为NumPy数组，以方便后续的处理与可视化。
• plt.subplot和plt.imshow:
  • 用于创建子图和展示图像，常用于数据可视化。

创建网络

当处理完数据后，就可以来进行网络的搭建了。网络搭建将逐一详细讨论生成器、判别器和损失函数。生成器G用到的是U-Net结构，输入的轮廓图编码再解码成真是图片，判别器D用到的是作者自己提出来的条件判别器PatchGAN，判别器D的作用是在轮廓图 的条件下，对于生成的图片判断为假，对于真实判断为真。

生成器G结构

U-Net是德国Freiburg大学模式识别和图像处理组提出的一种全卷积结构。它分为两个部分，其中左侧是由卷积和降采样操作组成的压缩路径，右侧是由卷积和上采样组成的扩张路径，扩张的每个网络块的输入由上一层上采样的特征和压缩路径部分的特征拼接而成。网络模型整体是一个U形的结构，因此被叫做U-Net。和常见的先降采样到低维度，再升采样到原始分辨率的编解码结构的网络相比，U-Net的区别是加入skip-connection，对应的feature maps和decode之后的同样大小的feature maps按通道拼一起，用来保留不同分辨率下像素级的细节信息。

定义UNet Skip Connection Block

import mindspore  # 导入MindSpore库
import mindspore.nn as nn  # 导入MindSpore的神经网络模块
import mindspore.ops as ops  # 导入MindSpore的操作模块

# 定义U-Net跳跃连接块
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
        
        # 根据norm_mode决定归一化方式
        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)  # 下采样ReLU激活
        up_relu = nn.ReLU()  # 上采样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()]  # 上采样部分，最后使用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))  # 添加Dropout层以防止过拟合

        self.model = nn.SequentialCell(model)  # 将所有层组合为一个SequentialCell
        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  # 返回输出

代码解析：

1. import mindspore:
   • 导入MindSpore库，以便使用其功能。
2. import mindspore.nn as nn:
   • 导入MindSpore的神经网络模块，命名为nn，用于构建神经网络。
3. import mindspore.ops as ops:
   • 导入MindSpore的操作模块，命名为ops，用于执行各种操作。
4. class UNetSkipConnectionBlock(nn.Cell)::
   • 定义一个UNetSkipConnectionBlock类，继承自nn.Cell，用于构建U-Net网络的跳跃连接块。
5. def __init__(self, ...)::
   • 构造函数，初始化网络层和参数。
   • outer_nc, inner_nc: 分别表示外层和内层的通道数。
   • in_planes: 输入通道数。
   • dropout: 是否使用Dropout。
   • submodule: 可能是下一个子模块。
   • outermost和innermost: 确定当前模块是外层还是内层。
   • alpha: LeakyReLU的负斜率。
   • norm_mode: 归一化模式（批归一化或实例归一化）。
6. down_norm和up_norm:
   • 根据选择的归一化模式，创建不同的归一化层。
7. down_conv和up_conv:
   • 定义下采样和上采样的卷积层，使用Conv2d和Conv2dTranspose。
8. down_relu和up_relu:
   • 定义激活函数，分别用于下采样和上采样。
9. 条件构造网络结构：
   • 根据outermost, innermost和其他参数，通过条件语句构建不同的网络结构。
10. self.model = nn.SequentialCell(model):
    • 将定义的层组合成一个SequentialCell以便于使用。
11. def construct(self, x)::
    • 实现前向传播方法，接收输入x并返回输出。
    • 如果使用跳跃连接，将输入x与输出out在通道维度上拼接。

API解析：

• nn.Cell:
  • MindSpore中的基本构建块，定义了神经网络的结构和前向传播逻辑。
• nn.Conv2d和nn.Conv2dTranspose:
  • 用于定义二维卷积和转置卷积（反卷积）层。
• nn.BatchNorm2d:
  • 用于进行二维批归一化，减少内部协变量偏移。
• nn.LeakyReLU和nn.ReLU:
  • 激活函数，用于加入非线性特征。
• nn.Dropout:
  • 用于随机丢弃神经元的输出，以防止过拟合。
• ops.concat:
  • 用于在指定维度上拼接张量，可以用于实现跳跃连接。
• nn.SequentialCell:
  • 用于将多个神经网络层组合成一个顺序执行的模型。

基于UNet的生成器

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跳跃连接块，使用ngf * 8的通道数
        unet_block = UNetSkipConnectionBlock(ngf * 8, ngf * 8, in_planes=None, submodule=None,
                                             norm_mode=norm_mode, innermost=True)

        # 通过循环创建多个UNet跳跃连接块，添加到unet_block中
        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跳跃连接块
        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)

        # 创建最外层的UNet跳跃连接块，输出通道为out_planes
        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)  # 前向传播，返回模型的输出

代码解析：

1. class UNetGenerator(nn.Cell)::
   • 定义一个UNetGenerator类，继承自nn.Cell，用于构建U-Net生成器。
2. def __init__(self, in_planes, out_planes, ngf=64, n_layers=8, norm_mode='bn', dropout=False)::
   • 构造函数，初始化生成器的参数。
   • in_planes: 输入通道数。
   • out_planes: 输出通道数。
   • ngf: 基础特征图的通道数，默认为64。
   • n_layers: 网络的层数，默认为8。
   • norm_mode: 归一化模式，默认为批归一化（‘bn’）。
   • dropout: 是否使用Dropout。
3. super(UNetGenerator, self).__init__():
   • 调用父类的构造函数，初始化基础类。
4. unet_block = UNetSkipConnectionBlock(...):
   • 创建最内层的U-Net跳跃连接块，使用ngf * 8的通道数，并设置innermost=True。
5. for _ in range(n_layers - 5)::
   • 通过循环创建多个U-Net跳跃连接块，构建网络的中间层，使用相同的通道配置。
6. 创建上层的UNet跳跃连接块：
   • unet_block = UNetSkipConnectionBlock(ngf * 4, ngf * 8, ...)
   • unet_block = UNetSkipConnectionBlock(ngf * 2, ngf * 4, ...)
   • unet_block = UNetSkipConnectionBlock(ngf, ngf * 2, ...)
   • 逐层搭建U-Net的结构，逐步减少特征图的通道数。
7. self.model = UNetSkipConnectionBlock(out_planes, ngf, in_planes=in_planes, submodule=unet_block, outermost=True, norm_mode=norm_mode):
   • 创建最外层的U-Net跳跃连接块，输出通道为out_planes，并设置outermost=True。
8. def construct(self, x)::
   • 定义前向传播方法，接收输入x并返回模型的输出。
9. return self.model(x):
   • 通过构建的模型进行前向传播，返回生成的输出。

API解析：

• nn.Cell:
  • MindSpore中的基本构建块，定义了神经网络的结构和前向传播逻辑。
• UNetSkipConnectionBlock:
  • 先前定义的跳跃连接块类，负责搭建U-Net的各个层。
• super():
  • 用于调用父类的方法，确保基类正确初始化。
• construct:
  • 定义了该类的前向传播逻辑，接收输入并计算输出。
• ngf, in_planes, out_planes, norm_mode, dropout:
  • 参数用于灵活配置网络结构和行为，支持多种配置，增强模型的可适应性。

原始cGAN的输入是条件x和噪声z两种信息，这里的生成器只使用了条件信息，因此不能生成多样性的结果。因此Pix2Pix在训练和测试时都使用了dropout，这样可以生成多样性的结果。

基于PatchGAN的判别器

判别器使用的PatchGAN结构，可看做卷积。生成的矩阵中的每个点代表原图的一小块区域（patch）。通过矩阵中的各个值来判断原图中对应每个Patch的真假。

import mindspore.nn as nn  # 导入MindSpore的神经网络模块

# 定义卷积-归一化-ReLU模块
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()  # 创建ReLU激活函数
            if alpha > 0:  # 如果使用LeakyReLU
                relu = nn.LeakyReLU(alpha)  # 创建LeakyReLU激活函数
            layers.append(relu)  # 将激活函数添加到层列表
        
        # 将所有层组合为一个SequentialCell
        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  # 初始化通道数

        # 创建多个卷积-归一化-ReLU模块
        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)  # 将所有层组合为一个SequentialCell

    def construct(self, x, y):
        x_y = ops.concat((x, y), axis=1)  # 在通道维度上拼接输入x和y
        output = self.features(x_y)  # 前向传播，计算输出
        return output  # 返回输出

代码解析：

1. import mindspore.nn as nn:
   • 导入MindSpore中的神经网络模块，为后续构建网络提供基础。
2. class ConvNormRelu(nn.Cell)::
   • 定义一个ConvNormRelu类，继承自nn.Cell，用于构建卷积-归一化-ReLU模块。
3. def __init__(self, ...)::
   • 构造函数，初始化卷积层、归一化层和激活函数的参数。
   • in_planes: 输入通道数。
   • out_planes: 输出通道数。
   • kernel_size: 卷积核大小。
   • stride: 卷积步长。
   • alpha: LeakyReLU的负斜率。
   • norm_mode: 归一化模式（‘batch’或’instance’）。
   • pad_mode: 填充模式。
   • use_relu: 是否使用ReLU激活函数。
   • padding: 填充大小。
4. super(ConvNormRelu, self).__init__():
   • 调用父类构造函数，初始化基础类。
5. norm = nn.BatchNorm2d(out_planes):
   • 创建批归一化层。
6. if norm_mode == 'instance'::
   • 根据选择的归一化模式，创建实例归一化层。
7. conv = nn.Conv2d(...):
   • 定义卷积层，设置通道数、卷积核大小、步长等参数。
8. layers.append(relu):
   • 根据选择的激活函数，将其添加到层列表中。
9. self.features = nn.SequentialCell(layers):
   • 将所有层组合成一个SequentialCell，便于使用。
10. def construct(self, x)::
    • 定义前向传播方法，计算输出。
11. class Discriminator(nn.Cell)::
    • 定义判别器网络。
12. def __init__(self, ...)::
    • 构造函数，初始化判别器的参数。
    • in_planes: 输入图像的通道数。
    • ndf: 基础特征图的通道数，默认为64。
    • n_layers: 网络的层数，默认为3。
13. layers.append(ConvNormRelu(...)):
    • 添加多个卷积-归一化-ReLU模块，以构建判别器的特征提取部分。
14. self.features = nn.SequentialCell(layers):
    • 将所有层组合为一个SequentialCell，形成完整的判别器网络。
15. def construct(self, x, y)::
    • 定义前向传播方法，接收两个输入x和y。
16. x_y = ops.concat((x, y), axis=1):
    • 在通道维度上拼接输入x和y，形成联合输入。
17. output = self.features(x_y):
    • 前向传播，计算输出。

API解析：

• nn.Cell:
  • MindSpore中的基本构建块，定义神经网络的结构和前向传播逻辑。
• nn.Conv2d:
  • 用于二维卷积操作的类。
• nn.BatchNorm2d和nn.LeakyReLU:
  • 分别用于二维批归一化和带泄漏的ReLU激活函数。
• nn.Pad:
  • 用于对输入进行填充。
• ops.concat:
  • 用于在指定维度上拼接多个张量。
• nn.SequentialCell:
  • 用于将多个层组合成一个顺序执行的模型。

Pix2Pix的生成器和判别器初始化

实例化Pix2Pix生成器和判别器。、

import mindspore.nn as nn  # 导入MindSpore的神经网络模块
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  # LeakyReLU的负斜率
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':  # 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))  # 初始化gamma为1
        cell.beta.set_data(init.initializer('zeros', cell.beta.shape))  # 初始化beta为0

# 创建判别器网络
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':  # 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))  # 初始化gamma为1
        cell.beta.set_data(init.initializer('zeros', cell.beta.shape))  # 初始化beta为0

# 定义Pix2Pix模型网络
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  # 返回生成的假图像

代码解析：

1. import mindspore.nn as nn:
   • 导入MindSpore的神经网络模块，用于构建深度学习模型。
2. from mindspore.common import initializer as init:
   • 导入MindSpore的初始化器模块，用于初始化模型参数。
3. 定义网络参数：
   • g_in_planes, g_out_planes, g_ngf, g_layers：生成器的输入、输出通道数、基础特征图通道数和层数。
   • d_in_planes, d_ndf, d_layers: 判别器的输入通道数、基础特征图通道数和层数。
   • alpha, init_gain, init_type: 用于初始化和激活函数的设置。
4. net_generator = UNetGenerator(...):
   • 创建生成器网络实例，使用指定的参数。
5. 初始化生成器中的卷积层和归一化层：
   • 使用cells_and_names()遍历生成器的所有子层。
   • 根据层的类型（卷积层或批归一化层）应用不同的初始化方法。
6. net_discriminator = Discriminator(...):
   • 创建判别器网络实例，使用指定的参数。
7. 初始化判别器中的卷积层和归一化层：
   • 同样使用cells_and_names()遍历判别器的所有子层，进行参数初始化。
8. 定义Pix2Pix模型：
   • class Pix2Pix(nn.Cell):：定义Pix2Pix网络，继承自nn.Cell。
   • def __init__(self, discriminator, generator):：构造函数，接受判别器和生成器作为参数。
   • super(Pix2Pix, self).__init__(auto_prefix=True)：调用父类构造函数，启用自动前缀。
   • self.net_discriminator = discriminator和self.net_generator = generator：保存判别器和生成器。
9. def construct(self, reala):：
   • 定义前向传播方法，接收真实图像reala作为输入。
   • fakeb = self.net_generator(reala)：通过生成器生成假图像。
   • return fakeb：返回生成的假图像。

API解析：

• nn.Cell：
  • MindSpore中的基本构建块，用于定义神经网络的结构和前向传播逻辑。
• init.initializer(...)：
  • 用于初始化网络权重的各种初始化方法（正态分布、Xavier均匀分布等）。
• nn.Conv2d和nn.BatchNorm2d：
  • 经典的卷积层和批归一化层，分别用于特征提取和归一化。
• cells_and_names()：
  • 用于遍历模型中的所有子单元及其名称，可以用于参数初始化等操作。
• auto_prefix=True：
  • 使得在调用父类时自动为每个层添加前缀，便于层的管理。

训练

训练分为两个主要部分：训练判别器和训练生成器。训练判别器的目的是最大程度地提高判别图像真伪的概率。训练生成器是希望能产生更好的虚假图像。在这两个部分中，分别获取训练过程中的损失，并在每个周期结束时进行统计。
下面进行训练：

import numpy as np  # 导入NumPy，用于数值计算
import os  # 导入os模块，用于文件和目录操作
import datetime  # 导入datetime模块，用于时间操作
from mindspore import value_and_grad, Tensor  # 从MindSpore导入必要的函数和类

# 定义训练参数
epoch_num = 100  # 总的训练轮数
ckpt_dir = "results/ckpt"  # 检查点保存目录
dataset_size = 400  # 数据集大小
val_pic_size = 256  # 验证图片大小
lr = 0.0002  # 初始学习率
n_epochs = 100  # 训练轮数
n_epochs_decay = 100  # 衰减轮数

def get_lr():
    """
    获取学习率
    返回一个包含每个训练步的学习率的Tensor
    """
    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))  # 返回学习率Tensor

# 创建数据集
dataset = ds.MindDataset("./dataset/dataset_pix2pix/train.mindrecord", columns_list=["input_images", "target_images"], shuffle=True, num_parallel_workers=16)
steps_per_epoch = dataset.get_dataset_size()  # 获取每个epoch的步数
loss_f = nn.BCEWithLogitsLoss()  # 创建二元交叉熵损失函数
l1_loss = nn.L1Loss()  # 创建L1损失函数

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  # GAN损失权重
    lambda_l1 = 100  # L1损失权重
    fakeb = net_generator(reala)  # 生成假图像
    pred0 = net_discriminator(reala, fakeb)  # 判别假图像
    loss_1 = loss_f(pred0, ops.ones_like(pred0))  # GAN损失
    loss_2 = l1_loss(fakeb, realb)  # L1损失
    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"])  # 输入图像Tensor
        target_image = Tensor(data["target_images"])  # 目标图像Tensor
        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")  # 保存生成器的模型参数

代码解析：

1. 导入模块：
   • import numpy as np, import os, import datetime: 导入必要的库和模块。
   • from mindspore import value_and_grad, Tensor: 从MindSpore导入用于计算梯度和Tensor的功能。
2. 定义训练参数：
   • 设置训练轮数、检查点目录、数据集大小、学习率、衰减轮数等。
3. 学习率获取函数：
   • get_lr(): 生成每一步的学习率，开始时保持为初始学习率，随着衰减轮数逐渐降低。
4. 创建数据集：
   • dataset = ds.MindDataset(...): 加载训练数据集，指定输入和目标图像列。
5. 损失函数定义：
   • loss_f = nn.BCEWithLogitsLoss(): 定义二元交叉熵损失函数。
   • l1_loss = nn.L1Loss(): 定义L1损失函数。
6. 判别器的前向传播：
   • forword_dis(reala, realb): 计算判别器的损失，包括对真实和生成图像的判别。
7. 生成器的前向传播：
   • forword_gan(reala, realb): 计算生成器的损失，包括GAN损失和L1损失。
8. 优化器创建：
   • 使用Adam优化器分别为判别器和生成器创建学习参数。
9. 计算梯度：
   • 使用value_and_grad计算损失的梯度。
10. 训练步骤：
    • train_step(reala, realb): 计算损失并更新网络参数。
11. 创建检查点目录：
    • 检查指定的目录是否存在，如果不存在就创建。
12. 训练过程：
    • 遍历每个epoch，并在每个数据批次上执行训练步骤，输出损失和运行时间。
    • 在训练结束时保存生成器的模型参数。

API解析：

• Tensor:
  • MindSpore中的基本数据结构，表示多维数组。
• value_and_grad:
  • 用于计算函数的输出和相应的梯度的函数。
• nn.BCEWithLogitsLoss:
  • 二元交叉熵损失函数，通常用于二分类任务。
• nn.L1Loss:
  • 计算输入与目标间的L1损失，即绝对误差。
• ds.MindDataset:
  • MindSpore的数据集类，用于加载和处理数据。
• mindspore.save_checkpoint:
  • 保存模型的参数到指定的路径。

推理

获取上述训练过程完成后的ckpt文件，通过load_checkpoint和load_param_into_net将ckpt中的权重参数导入到模型中，获取数据进行推理并对推理的效果图进行演示（由于时间问题，训练过程只进行了100个epoch）。

from mindspore import load_checkpoint, load_param_into_net  # 导入MindSpore的检查点加载和参数设置功能

# 加载生成器的模型参数
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):  # 显示前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()  # 显示图像

代码解析：

1. 导入模块：
   • from mindspore import load_checkpoint, load_param_into_net：
     • 导入用于加载模型参数的函数。
2. 加载模型参数：
   • param_g = load_checkpoint(ckpt_dir + "Generator.ckpt")：
     • 从检查点文件中加载生成器的参数。
   • load_param_into_net(net_generator, param_g)：
     • 将加载的参数导入到生成器网络中，确保网络结构与参数一致。
3. 创建数据集和加载数据：
   • dataset = ds.MindDataset(...)：
     • 创建MindSpore数据集，指定输入和目标图像列，设置为随机打乱。
   • data_iter = next(dataset.create_dict_iterator())：
     • 使用数据集创建一个字典迭代器，并获取一个批次的数据。
4. 生成预测图像：
   • predict_show = net_generator(data_iter["input_images"])：
     • 将输入图像传入生成器，得到生成的图像。
5. 绘制输入和生成图像：
   • plt.figure(figsize=(10, 3), dpi=140)：
     • 创建一个新的图形，设置大小和分辨率。
   • for i in range(10):：
     • 遍历前10个图像，创建子图。
   • plt.subplot(2, 10, i + 1)：
     • 创建一个2行10列的子图布局，显示输入图像。
   • plt.imshow(...)：
     • 显示输入图像，并进行必要的格式转换以适应显示标准（将张量转换为NumPy数组并调整维度）。
   • plt.axis("off")：
     • 关闭子图的坐标轴显示。
   • plt.subplots_adjust(...)：
     • 调整子图之间的空隙。
   • plt.show()：
     • 显示整个图形。

API解析：

• load_checkpoint：
  • 用于加载模型保存的参数，返回一个包含所有参数的字典。
• load_param_into_net：
  • 将加载的参数字典导入到指定的网络中，确保网络可以使用这些参数进行推理或训练。
• ds.MindDataset：
  • MindSpore的数据集类，用于加载和处理数据。
• Tensor.asnumpy()：
  • 将MindSpore的Tensor转换为NumPy数组。
• plt.imshow(...)：
  • Matplotlib用于显示图像的函数。
• plt.subplot(...)：
  • 用于创建子图的函数，可指定子图的位置和布局。