风格迁移2-07：MUNIT(多模态无监督)-源码无死角解析（3）-loss计算

以下链接是个人关于 MUNIT(多模态无监督)-图片风格转换，的所有见解，如有错误欢迎大家指出，我会第一时间纠正。有兴趣的朋友可以加微信 17575010159 相互讨论技术。若是帮助到了你什么，一定要记得点赞！因为这是对我最大的鼓励。$文末附带$$公众号-$$海量资源。$

风格迁移2-00：MUNIT(多模态无监督)-目录-史上最新无死角讲解

前言

通过上一篇博客的介绍，我们已经知道了网络前向传播的过程，主要代码再 trainer.py 中实现，现在我们来看看网络的是如何计算 loss 进行优化的。其实现，主要包含了两个函数，如下：

    # 生成模型进行优化
    def gen_update(self, x_a, x_b, hyperparameters):

    # 鉴别模型进行优化
    def dis_update(self, x_a, x_b, hyperparameters):

gen_update

该函数代码的实现如下：

    def gen_update(self, x_a, x_b, hyperparameters):
        # 给输入的 x_a, x_b 加入随机噪声
        self.gen_opt.zero_grad()
        s_a = Variable(torch.randn(x_a.size(0), self.style_dim, 1, 1).cuda())
        s_b = Variable(torch.randn(x_b.size(0), self.style_dim, 1, 1).cuda())

        # 对 a, b 图像进行解码在编码（自我解码，自我编码，没有进行风格交换），
        c_a, s_a_prime = self.gen_a.encode(x_a)
        c_b, s_b_prime = self.gen_b.encode(x_b)
        # decode (within domain)，
        x_a_recon = self.gen_a.decode(c_a, s_a_prime)
        x_b_recon = self.gen_b.decode(c_b, s_b_prime)


        # 进行交叉解码，即两张图片的content code，style code进行互换
        # decode (cross domain)，
        x_ba = self.gen_a.decode(c_b, s_a)
        x_ab = self.gen_b.decode(c_a, s_b)

        # encode again，对上面合成的图片再进行编码，得到重构的content code，style code
        c_b_recon, s_a_recon = self.gen_a.encode(x_ba)
        c_a_recon, s_b_recon = self.gen_b.encode(x_ab)


        # decode again (if needed)，重构的content code 与真实图片编码得到 style code（s_x_prime）进行解码，生成新图片
        x_aba = self.gen_a.decode(c_a_recon, s_a_prime) if hyperparameters['recon_x_cyc_w'] > 0 else None
        x_bab = self.gen_b.decode(c_b_recon, s_b_prime) if hyperparameters['recon_x_cyc_w'] > 0 else None



        # reconstruction loss，计算重构的loss
        # 重构图片a，与真实图片a计算计算loss
        self.loss_gen_recon_x_a = self.recon_criterion(x_a_recon, x_a)
        # 重构图片b，与真实图片b计算计算loss
        self.loss_gen_recon_x_b = self.recon_criterion(x_b_recon, x_b)
        # 合成图片，再编码得到的style code，与正太分布的随机生成的s_x（style code）计算loss
        self.loss_gen_recon_s_a = self.recon_criterion(s_a_recon, s_a)
        self.loss_gen_recon_s_b = self.recon_criterion(s_b_recon, s_b)

        # 由合成图片编码得到的content code，与真实的图片编码得到的content code 计算loss
        self.loss_gen_recon_c_a = self.recon_criterion(c_a_recon, c_a)
        self.loss_gen_recon_c_b = self.recon_criterion(c_b_recon, c_b)

        # 循环一致性loss，使用:
        # 1、通过对合成图片(x_ab)进行编码得到的content code，
        # 2、对真实图片(x_a)进行编码，得到style code
        # 3、结合content code，style code 进行解码，得到 x_aba，然后计算x_aba 与 x_a 的loss
        self.loss_gen_cycrecon_x_a = self.recon_criterion(x_aba, x_a) if hyperparameters['recon_x_cyc_w'] > 0 else 0
        self.loss_gen_cycrecon_x_b = self.recon_criterion(x_bab, x_b) if hyperparameters['recon_x_cyc_w'] > 0 else 0

        # GAN loss，最终生成图片与真实图片之间的loss
        self.loss_gen_adv_a = self.dis_a.calc_gen_loss(x_ba)
        self.loss_gen_adv_b = self.dis_b.calc_gen_loss(x_ab)


        # domain-invariant perceptual loss，使用VGG计算感知loss
        self.loss_gen_vgg_a = self.compute_vgg_loss(self.vgg, x_ba, x_b) if hyperparameters['vgg_w'] > 0 else 0
        self.loss_gen_vgg_b = self.compute_vgg_loss(self.vgg, x_ab, x_a) if hyperparameters['vgg_w'] > 0 else 0

        # 全局loss
        # total loss
        self.loss_gen_total = hyperparameters['gan_w'] * self.loss_gen_adv_a + \
                              hyperparameters['gan_w'] * self.loss_gen_adv_b + \
                              hyperparameters['recon_x_w'] * self.loss_gen_recon_x_a + \
                              hyperparameters['recon_s_w'] * self.loss_gen_recon_s_a + \
                              hyperparameters['recon_c_w'] * self.loss_gen_recon_c_a + \
                              hyperparameters['recon_x_w'] * self.loss_gen_recon_x_b + \
                              hyperparameters['recon_s_w'] * self.loss_gen_recon_s_b + \
                              hyperparameters['recon_c_w'] * self.loss_gen_recon_c_b + \
                              hyperparameters['recon_x_cyc_w'] * self.loss_gen_cycrecon_x_a + \
                              hyperparameters['recon_x_cyc_w'] * self.loss_gen_cycrecon_x_b + \
                              hyperparameters['vgg_w'] * self.loss_gen_vgg_a + \
                              hyperparameters['vgg_w'] * self.loss_gen_vgg_b

        # 反向传播
        self.loss_gen_total.backward()
        self.gen_opt.step()

以上的注释还是比较详细的，主要的 loss 分为4个部分：

1. 重构图片与真实图片之间的 loss
2. 重构图片编码得到的 latent code 与 真实图片编码得到的latent code 计算loss
3. 图片翻译到目标域，再反回来和原图计算 loss
4. 使用VGG计算域感知 loss

dis_update

    # 鉴别模型进行优化
    def dis_update(self, x_a, x_b, hyperparameters):
        self.dis_opt.zero_grad()
        # 随机生成符合正太分布的style code
        s_a = Variable(torch.randn(x_a.size(0), self.style_dim, 1, 1).cuda())
        s_b = Variable(torch.randn(x_b.size(0), self.style_dim, 1, 1).cuda())

        # encode，对输入的图片进行编码，得到 content code 以及 style code
        c_a, _ = self.gen_a.encode(x_a)
        c_b, _ = self.gen_b.encode(x_b)

        # 交叉进行解码（即互换 content code 或者 style code）
        # decode (cross domain)
        x_ba = self.gen_a.decode(c_b, s_a)
        x_ab = self.gen_b.decode(c_a, s_b)

        # D loss 计算鉴别器的loss
        self.loss_dis_a = self.dis_a.calc_dis_loss(x_ba.detach(), x_a)
        self.loss_dis_b = self.dis_b.calc_dis_loss(x_ab.detach(), x_b)
        self.loss_dis_total = hyperparameters['gan_w'] * self.loss_dis_a + hyperparameters['gan_w'] * self.loss_dis_b
        self.loss_dis_total.backward()
        self.dis_opt.step()

鉴别器的 loss 相对来说是比较简单的，就是对生成的图片进行鉴别。知道生成模型去学习，生成逼真的图像。

sample

下面再介绍一下 def sample(self, x_a, x_b) 函数，如下：

    def sample(self, x_a, x_b):
        self.eval()
        # 当前训练过程中使用的 style code
        s_a1 = Variable(self.s_a)
        s_b1 = Variable(self.s_b)
        # 零时随机生成的 style code(符合正太分布)
        s_a2 = Variable(torch.randn(x_a.size(0), self.style_dim, 1, 1).cuda())
        s_b2 = Variable(torch.randn(x_b.size(0), self.style_dim, 1, 1).cuda())
        x_a_recon, x_b_recon, x_ba1, x_ba2, x_ab1, x_ab2 = [], [], [], [], [], []


        for i in range(x_a.size(0)):
            # 输入图片a,b，分别得到对应的content code 以及 style code
            c_a, s_a_fake = self.gen_a.encode(x_a[i].unsqueeze(0))
            c_b, s_b_fake = self.gen_b.encode(x_b[i].unsqueeze(0))

            # 对图片a进行重构，使用从图片a分离出来的content code 以及 style code
            x_a_recon.append(self.gen_a.decode(c_a, s_a_fake))
            # 对图片b进行重构，使用从图片b分离出来的content code 以及 style code
            x_b_recon.append(self.gen_b.decode(c_b, s_b_fake))

            # 使用分离出来的content code， 结合符合正太分布随机生成的style code，生成图片
            x_ba1.append(self.gen_a.decode(c_b, s_a1[i].unsqueeze(0)))
            x_ba2.append(self.gen_a.decode(c_b, s_a2[i].unsqueeze(0)))
            x_ab1.append(self.gen_b.decode(c_a, s_b1[i].unsqueeze(0)))
            x_ab2.append(self.gen_b.decode(c_a, s_b2[i].unsqueeze(0)))

        #把图片的像素连接起来
        x_a_recon, x_b_recon = torch.cat(x_a_recon), torch.cat(x_b_recon)
        x_ba1, x_ba2 = torch.cat(x_ba1), torch.cat(x_ba2)
        x_ab1, x_ab2 = torch.cat(x_ab1), torch.cat(x_ab2)
        self.train()
        return x_a, x_a_recon, x_ab1, x_ab2, x_b, x_b_recon, x_ba1, x_ba2

该函数的作用，是为了再训练的时候，查看目前的效果，其在 trainer.py 中，可以看到被调用过程如下：

    # Write images，到达指定次数后，把生成的样本图片写入到输出文件夹，方便观察生成效果，重新保存
    if (iterations + 1) % config['image_save_iter'] == 0:
        with torch.no_grad():
            test_image_outputs = trainer.sample(test_display_images_a, test_display_images_b)
            train_image_outputs = trainer.sample(train_display_images_a, train_display_images_b)
        write_2images(test_image_outputs, display_size, image_directory, 'test_%08d' % (iterations + 1))
        write_2images(train_image_outputs, display_size, image_directory, 'train_%08d' % (iterations + 1))
        # HTML
        write_html(output_directory + "/index.html", iterations + 1, config['image_save_iter'], 'images')

    # Write images，到达指定次数后，把生成的样本图片写入到输出文件夹，方便观察生成效果，覆盖上一次结果
    if (iterations + 1) % config['image_display_iter'] == 0:
        with torch.no_grad():
            image_outputs = trainer.sample(train_display_images_a, train_display_images_b)
        write_2images(image_outputs, display_size, image_directory, 'train_current')

单训练到指定次数之后，就会把事先挑选出来的样本，通过 def sample(self, x_a, x_b) 函数进行推断，把推断之后的结果保存到 outputs 文件夹。