风格迁移2-06：MUNIT(多模态无监督)-源码无死角解析（3）-模型框架（前向传播）

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

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

前言

根据上一篇博客，可以知道，模型的构建代码为 train.py 中的如下部分：

    # Setup model and data loader, 根据配置创建模型
    if opts.trainer == 'MUNIT':
        trainer = MUNIT_Trainer(config)
    elif opts.trainer == 'UNIT':
        trainer = UNIT_Trainer(config)
    else:
        sys.exit("Only support MUNIT|UNIT")
    trainer.cuda()

那么我们就进入 MUNIT_Trainer(config) 看看把，可以知道MUNIT_Trainer主要实现了如下函数：

class MUNIT_Trainer(nn.Module):
	def __init__(self, hyperparameters):
	
    # 网络前向传播
    def forward(self, x_a, x_b):

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

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

以上几个函数，就是我们重点分析的对象。

框架总览

对于 pytorch 框架构建的网络，我们一般是先查看他的 forward(网络前向传播) ,代码实现如下：

    # 网络前向传播
    def forward(self, x_a, x_b):
        # 先设定为推断模式
        self.eval()
        # 把随机噪声转化为pytorch变量
        s_a = Variable(self.s_a)
        s_b = Variable(self.s_b)

        # 输入图片a，b进行编码，分别得到论文中的content code 以及 style code
        c_a, s_a_fake = self.gen_a.encode(x_a)
        c_b, s_b_fake = self.gen_b.encode(x_b)

        # 对content code 加入噪声，然后进行解码（混合），得到合成的图片
        x_ba = self.gen_a.decode(c_b, s_a)
        x_ab = self.gen_b.decode(c_a, s_b)
        self.train()
        return x_ab, x_ba

从总体来看，过程是十分简单的，首先对输入的两张图片都进行编码，分别得到两张图片的 content code 以及 style code，再互换，然后加入符合正态分布的噪声生成新的图片x_ab, x_ba。

了解了总体框架之后，我们再来看看初始化函数。

初始化函数

该函数的注释如下(后面有带读)：

    def __init__(self, hyperparameters):
        super(MUNIT_Trainer, self).__init__()
        lr = hyperparameters['lr']
        # Initiate the networks
        # 生成网络模型a, 即由数据集A到数据集B的映射
        self.gen_a = AdaINGen(hyperparameters['input_dim_a'], hyperparameters['gen'])  # auto-encoder for domain a
        # 生成网络模型b, 即由数据集B到数据集A的映射
        self.gen_b = AdaINGen(hyperparameters['input_dim_b'], hyperparameters['gen'])  # auto-encoder for domain b
        # 鉴别模型a，鉴别生成的图像，是否和数据集A的分布一致
        self.dis_a = MsImageDis(hyperparameters['input_dim_a'], hyperparameters['dis'])  # discriminator for domain a
        # 鉴别模型b，鉴别生成的图像，是否和数据集B的分布一致
        self.dis_b = MsImageDis(hyperparameters['input_dim_b'], hyperparameters['dis'])  # discriminator for domain b
        # 使用正则化的方式
        self.instancenorm = nn.InstanceNorm2d(512, affine=False)
        # style 输出的特征码维度
        self.style_dim = hyperparameters['gen']['style_dim']

        # fix the noise used in sampling， 随机加入噪声，噪声符合正态分布
        display_size = int(hyperparameters['display_size'])
        self.s_a = torch.randn(display_size, self.style_dim, 1, 1).cuda()
        self.s_b = torch.randn(display_size, self.style_dim, 1, 1).cuda()

        # Setup the optimizers， 优化器的超参数
        beta1 = hyperparameters['beta1']
        beta2 = hyperparameters['beta2']

        # 鉴别模型a,b的相关参数
        dis_params = list(self.dis_a.parameters()) + list(self.dis_b.parameters())
        # 生成模型a,b的相关参数
        gen_params = list(self.gen_a.parameters()) + list(self.gen_b.parameters())

        # 构建鉴别模型以及生成生成模型的优化器
        self.dis_opt = torch.optim.Adam([p for p in dis_params if p.requires_grad],
                                        lr=lr, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay'])
        self.gen_opt = torch.optim.Adam([p for p in gen_params if p.requires_grad],
                                        lr=lr, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay'])
        # 鉴别模型以及生成生成模型的学习率衰减策略
        self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters)
        self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters)

        # Network weight initialization，网络模型权重初始化
        self.apply(weights_init(hyperparameters['init']))
        self.dis_a.apply(weights_init('gaussian'))
        self.dis_b.apply(weights_init('gaussian'))

        # Load VGG model if needed，加载VGG模型，用来计算感知 loss
        if 'vgg_w' in hyperparameters.keys() and hyperparameters['vgg_w'] > 0:
            self.vgg = load_vgg16(hyperparameters['vgg_model_path'] + '/models')
            self.vgg.eval()
            for param in self.vgg.parameters():
                param.requires_grad = False

总的来说，初始化的过程中，主要构建了两个生成器gen_a，gen_b。以及两个鉴别器dis_a，dis_b。和对应的优化器。最后还创建了计算感知 loss 需要的VGG网络。

最主要的是，生成器gen_a，gen_b中包含了解码器和生成器，下篇博客我会对 loss 的计算进行讲解，需要大家继续观看。