风格迁移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 codec_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 Nonex_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计算计算lossself.loss_gen_recon_x_a = self.recon_criterion(x_a_recon, x_a)# 重构图片b，与真实图片b计算计算lossself.loss_gen_recon_x_b = self.recon_criterion(x_b_recon, x_b)# 合成图片，再编码得到的style code，与正太分布的随机生成的s_x（style code）计算lossself.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 计算lossself.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 的lossself.loss_gen_cycrecon_x_a = self.recon_criterion(x_aba, x_a) if hyperparameters['recon_x_cyc_w'] > 0 else 0self.loss_gen_cycrecon_x_b = self.recon_criterion(x_bab, x_b) if hyperparameters['recon_x_cyc_w'] > 0 else 0# GAN loss，最终生成图片与真实图片之间的lossself.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计算感知lossself.loss_gen_vgg_a = self.compute_vgg_loss(self.vgg, x_ba, x_b) if hyperparameters['vgg_w'] > 0 else 0self.loss_gen_vgg_b = self.compute_vgg_loss(self.vgg, x_ab, x_a) if hyperparameters['vgg_w'] > 0 else 0# 全局loss# total lossself.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 codes_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 codec_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 计算鉴别器的lossself.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_bself.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 codes_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 codec_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 codex_a_recon.append(self.gen_a.decode(c_a, s_a_fake))# 对图片b进行重构，使用从图片b分离出来的content code 以及 style codex_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))# HTMLwrite_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 文件夹。