以下链接是个人关于 MUNIT(多模态无监督)-图片风格转换,的所有见解,如有错误欢迎大家指出,我会第一时间纠正。有兴趣的朋友可以加微信 17575010159 相互讨论技术。若是帮助到了你什么,一定要记得点赞!因为这是对我最大的鼓励。 文末附带 \color{blue}{文末附带} 文末附带 公众号 − \color{blue}{公众号 -} 公众号− 海量资源。 \color{blue}{ 海量资源}。 海量资源。
风格迁移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 的计算进行讲解,需要大家继续观看。