GFPGAN源码分析—第八篇

2021SC@SDUSC

源码：

models\init.py

models\gfpgan_model.py

本篇主要分析init.py与models\gfpgan_model.py下的

class GFPGANModel(BaseModel) 类init(self, opt) 方法

目录

init.py

gfpgan_model.py

class GFPGANModel(BaseModel)

init(self, opt)

init_training_settings(self)

---

init.py

自动扫描和导入注册表的模型模块

#在models文件夹下扫描所有以 '_model.py' 结尾的文件
model_folder = osp.dirname(osp.abspath(__file__))
model_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(model_folder) if v.endswith('_model.py')]
# 导入所有模型模块
_model_modules = [importlib.import_module(f'gfpgan.models.{file_name}') for file_name in model_filenames]

那么实际上就是导入models文件夹下gfpgan_model.py文件，接下来我们来看一下

gfpgan_model.py

本文件中只包含GFPGANModel(BaseModel)一个类

创建了一个MODEL_REGISTRY对象，并在类定义的时候用装饰器装饰它，以装饰器的形式调用MODEL_REGISTRY类的register函数

@MODEL_REGISTRY.register()
class GFPGANModel(BaseModel):
    """GFPGAN model for <Towards real-world blind faces restoratin with generative facial prior>"""

class GFPGANModel(BaseModel)

基于生成性人脸先验信息的真实盲脸修复 的 GFPGAN 模型

init(self, opt)

简单看一下代码

super(GFPGANModel, self).__init__(opt)
self.idx = 0

# 网络定义
self.net_g = build_network(opt['network_g'])
self.net_g = self.model_to_device(self.net_g)
self.print_network(self.net_g)

# 读取预训练的模型
load_path = self.opt['path'].get('pretrain_network_g', None)
#如果路径不为空
if load_path is not None:
    param_key = self.opt['path'].get('param_key_g', 'params')
    self.load_network(self.net_g, load_path, self.opt['path'].get('strict_load_g', True), param_key)

self.log_size = int(math.log(self.opt['network_g']['out_size'], 2))

if self.is_train:
    self.init_training_settings()

在读取预训练的模型时，实际上就是从train_gfpgan_v1.yml配置文件中读取到相应的参数的数值与路径。

init_training_settings(self)

初始化训练设置

1.读取opt['train']

train_opt = self.opt['train']

2.定义net_d

#构建网络
self.net_d = build_network(self.opt['network_d'])
#将模型放到gpu（cuda）上
self.net_d = self.model_to_device(self.net_d)
self.print_network(self.net_d)
# 读取与训练好的模型
load_path = self.opt['path'].get('pretrain_network_d', None)
if load_path is not None:
    self.load_network(self.net_d, load_path, self.opt['path'].get('strict_load_d', True))

3.定义net_g

# net_g_ema 仅用于在一个GPU上测试并保存
# 不需要使用DistributedDataParallel进行包装
self.net_g_ema = build_network(self.opt['network_g']).to(self.device)
# 读取预训练模型
load_path = self.opt['path'].get('pretrain_network_g', None)
if load_path is not None:
    self.load_network(self.net_g_ema, load_path, self.opt['path'].get('strict_load_g', True), 'params_ema')
else:
    self.model_ema(0)  # copy net_g weight

self.net_g.train()
self.net_d.train()
self.net_g_ema.eval()

根据配置文件：net_g读取预训练模型为arcface_resnet18.pth

4.面部组件网络

if ('network_d_left_eye' in self.opt and 'network_d_right_eye' in self.opt and 'network_d_mouth' in self.opt):
            self.use_facial_disc = True
        else:
            self.use_facial_disc = False

        if self.use_facial_disc:
            # left eye
            self.net_d_left_eye = build_network(self.opt['network_d_left_eye'])
            self.net_d_left_eye = self.model_to_device(self.net_d_left_eye)
            self.print_network(self.net_d_left_eye)
            load_path = self.opt['path'].get('pretrain_network_d_left_eye')
            if load_path is not None:
                self.load_network(self.net_d_left_eye, load_path, True, 'params')
            # right eye
            self.net_d_right_eye = build_network(self.opt['network_d_right_eye'])
            self.net_d_right_eye = self.model_to_device(self.net_d_right_eye)
            self.print_network(self.net_d_right_eye)
            load_path = self.opt['path'].get('pretrain_network_d_right_eye')
            if load_path is not None:
                self.load_network(self.net_d_right_eye, load_path, True, 'params')
            # mouth
            self.net_d_mouth = build_network(self.opt['network_d_mouth'])
            self.net_d_mouth = self.model_to_device(self.net_d_mouth)
            self.print_network(self.net_d_mouth)
            load_path = self.opt['path'].get('pretrain_network_d_mouth')
            if load_path is not None:
                self.load_network(self.net_d_mouth, load_path, True, 'params')

            self.net_d_left_eye.train()
            self.net_d_right_eye.train()
            self.net_d_mouth.train()

            # ----------- 定义面部组件的 gan loss ----------- #
            self.cri_component = build_loss(train_opt['gan_component_opt']).to(self.device)

5.定义损失

if train_opt.get('pixel_opt'):
    self.cri_pix = build_loss(train_opt['pixel_opt']).to(self.device)
else:
    self.cri_pix = None

if train_opt.get('perceptual_opt'):
    self.cri_perceptual = build_loss(train_opt['perceptual_opt']).to(self.device)
else:
    self.cri_perceptual = None
    # pyramid loss, component style loss, identity loss 都使用L1损失
    self.cri_l1 = build_loss(train_opt['L1_opt']).to(self.device)

    # gan loss (wgan)
    self.cri_gan = build_loss(train_opt['gan_opt']).to(self.device)

6.identity loss的定义

if 'network_identity' in self.opt:
    self.use_identity = True
else:
    self.use_identity = False

if self.use_identity:
    # 定义 identity network
    self.network_identity = build_network(self.opt['network_identity'])
    self.network_identity = self.model_to_device(self.network_identity)
    self.print_network(self.network_identity)
    load_path = self.opt['path'].get('pretrain_network_identity')
    if load_path is not None:
        self.load_network(self.network_identity, load_path, True, None)
    self.network_identity.eval()
    for param in self.network_identity.parameters():
        param.requires_grad = False

# 正则化权重
self.r1_reg_weight = train_opt['r1_reg_weight']  # for discriminator
self.net_d_iters = train_opt.get('net_d_iters', 1)
self.net_d_init_iters = train_opt.get('net_d_init_iters', 0)
self.net_d_reg_every = train_opt['net_d_reg_every']

# 设置优化器和调度程序
self.setup_optimizers()
self.setup_schedulers()