[人工智能]pixel2style2pixel（pSp）实现解读【二】 -- 代码层面

前言

pixel2stylepixel是人脸编辑中比较重要的encoder方法，这里针对代码部分作简要记录。
理论部分见here
论文
Code

项目文件介绍

主要关注encoder部分的代码设计，所以介绍下项目中的文件：

（1）models文件夹下psp.py

其实这就是psp的模型定义文件了，里面就包含了一个class类，后续的使用主要就是调用该文件夹下pSp类

（2）encoders文件夹下helpers.py

在这个文件夹下主要定义了bottlenck模型以及几个它的变种 【用于将真实图像生成不同维度的向量】

（3）encoders文件夹下psp_encoder.py

psp的encoder实现文件

代码解读

（一）encoder

自上而下来阅读

首先在psp.py文件中

#psp.py

self.opts.n_styles = int(math.log(self.opts.output_size, 2)) * 2 - 2

在这里通过最终输出的分辨率来计算style模块的层数，比如最终输出1024 * 1024图像，那n_styles便是18层 【这里是在style Mixing模块用到】

#psp.py
    
def set_encoder(self):
	if self.opts.encoder_type == 'GradualStyleEncoder':
		encoder = psp_encoders.GradualStyleEncoder(50, 'ir_se', self.opts)
	elif self.opts.encoder_type == 'BackboneEncoderUsingLastLayerIntoW':
		encoder = psp_encoders.BackboneEncoderUsingLastLayerIntoW(50, 'ir_se', self.opts)
	elif self.opts.encoder_type == 'BackboneEncoderUsingLastLayerIntoWPlus':
		encoder = psp_encoders.BackboneEncoderUsingLastLayerIntoWPlus(50, 'ir_se', self.opts)
	else:
		raise Exception('{} is not a valid encoders'.format(self.opts.encoder_type))
	return encoder

默认都是使用 ‘GradualStyleEncoder’模式，其他两种属于简单粗暴的方式。

其中关于ir，论文中说是 ResNet-IR，pretrained on face recognition，代码借鉴了该项目

ir_se即在ir的网络模型最后加上了se模型。

接下来我们便开始了解psp_encoders.py文件下的GradualStyleEncoder类了：

# psp_encoder.py

blocks = get_blocks(num_layers)

根据跳转，可以查看到：

#helpers.py

class Bottleneck(namedtuple('Block', ['in_channel', 'depth', 'stride'])):
	""" A named tuple describing a ResNet block. """


def get_block(in_channel, depth, num_units, stride=2):
	return [Bottleneck(in_channel, depth, stride)] + [Bottleneck(depth, depth, 1) for i in range(num_units - 1)]


def get_blocks(num_layers):
	if num_layers == 50:
		blocks = [
			get_block(in_channel=64, depth=64, num_units=3),
			get_block(in_channel=64, depth=128, num_units=4),
			get_block(in_channel=128, depth=256, num_units=14),
			get_block(in_channel=256, depth=512, num_units=3)
		]
	elif num_layers == 100:
		blocks = [
			get_block(in_channel=64, depth=64, num_units=3),
			get_block(in_channel=64, depth=128, num_units=13),
			get_block(in_channel=128, depth=256, num_units=30),
			get_block(in_channel=256, depth=512, num_units=3)
		]
	elif num_layers == 152:
		blocks = [
			get_block(in_channel=64, depth=64, num_units=3),
			get_block(in_channel=64, depth=128, num_units=8),
			get_block(in_channel=128, depth=256, num_units=36),
			get_block(in_channel=256, depth=512, num_units=3)
		]
	else:
		raise ValueError("Invalid number of layers: {}. Must be one of [50, 100, 152]".format(num_layers))
	return blocks

所以如果根据初始的输入num_layers == 50，那么会得到24层block（Bottleneck）

#psp_encoder.py

'''
unit_module为bottleneck_IR或bottleneck_IR_SE
'''

modules = []
        for block in blocks:
            for bottleneck in block:
                modules.append(unit_module(bottleneck.in_channel,
                                           bottleneck.depth,
                                           bottleneck.stride))
self.body = Sequential(*modules)

#psp_encoders.py

self.style_count = opts.n_styles
self.coarse_ind = 3
self.middle_ind = 7
for i in range(self.style_count):
	if i < self.coarse_ind:
		style = GradualStyleBlock(512, 512, 16)
     elif i < self.middle_ind:
         style = GradualStyleBlock(512, 512, 32)
     else:
         style = GradualStyleBlock(512, 512, 64)

以上便是对于FPN结构的实现，对于不同层次（coarse，middle，fine），变换的只是spatial参数，而该参数影响的只是网络模型里卷积层的个数，所以在低层次中，卷积层个数少，学习到的特征少，高层次中，卷积层个数多，学习到的特征也会多。

#psp_encoders.py

modulelist = list(self.body._modules.values())
for i, l in enumerate(modulelist):
	x = l(x)
	if i == 6:
		c1 = x
	elif i == 20:
		c2 = x
	elif i == 23:
		c3 = x

从指定 layer （第 6，20 和 23 层）拿中间的特征图，其实回看之前的block可以发现这几个层刚好属于分界点，可以认为是FPN结构中coarse、mid、fine的分界点，输入x的维度分别为128，256，512。

而下面的操作便是向list添加latent code，最后将其送入w+空间了。【在这个过程中加入了_upsample_add，用于上采样+合并两个特征图】

#psp.py

'''
forward函数的逻辑
'''
if alpha is not None:
    codes[:, i] = alpha * inject_latent[:, i] + (1 - alpha) * codes[:, i]
else:
    codes[:, i] = inject_latent[:, i]

首先把输入图片直接扔到 encoder 里拿到 latent code。

如果选择了 start_from_latent_avg，则 code 加上之前的平均 latent code。

如果要用 style mixing，这里则借助一个 latent_mask 和 inject_latent mix 起来

（二）关于数据transform

1. 图像 encode 任务：input 图像和 target 图像相同，应用了 RandomHorizontalFlip(0.5) 来做数据增强。
2. 人脸正面化任务：input 图像和 target 各自 RandomHorizontalFlip(0.5)。
3. SketchToImage 任务：input 图像没有做 Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]) 的处理，而目标图像做了。
4. SegToImage 任务：这里有个 ToOneHot 的 transform 我没看懂，无关紧要，算了。
5. 超分辨任务：这里对 input 图像做了随机倍数的下采样，之后 resize 回 256。
6. 上色任务：这里是我自己加的，只需要对 input 图像做一下 transforms.Grayscale(num_output_channels=3) 就好。

参考

pixel2style2pixel（pSp）代码实现解读