Add Fom mobile train code (PaddlePaddle#411)

* add fom mobile training

configs/firstorder_vox_mobile_256.yaml

@@ -4,8 +4,8 @@ output_dir: output_dir
 dataset:
   train:
     name: FirstOrderDataset
-    batch_size: 1
-    num_workers: 1
+    batch_size: 8
+    num_workers: 4
     use_shared_memory: False
     phase: train
     dataroot: data/first_order/Voxceleb/
@@ -38,7 +38,10 @@ dataset:
 
 
 model:
-  name: FirstOrderModel
+  name: FirstOrderModelMobile
+  mode: generator # should be kp_detector, generator, both
+  kp_weight_path: None
+  gen_weight_path: None
   common_params:
     num_kp: 10
     num_channels: 3
@@ -64,6 +67,25 @@ model:
         num_blocks: 5
         scale_factor: 0.25
       mobile_net: True
+  generator_ori:
+    name: FirstOrderGenerator
+    kp_detector_cfg:
+      temperature: 0.1
+      block_expansion: 32
+      max_features: 1024
+      scale_factor: 0.25
+      num_blocks: 5
+    generator_cfg:
+      block_expansion: 64
+      max_features: 512
+      num_down_blocks: 2
+      num_bottleneck_blocks: 6
+      estimate_occlusion_map: True
+      dense_motion_params:
+        block_expansion: 64
+        max_features: 1024
+        num_blocks: 5
+        scale_factor: 0.25
   discriminator:
     name: FirstOrderDiscriminator
     discriminator_cfg:

docs/zh_CN/tutorials/motion_driving.md

@@ -103,6 +103,7 @@ export CUDA_VISIBLE_DEVICES=0
 python tools/main.py --config-file configs/dataset_name.yaml
 ```
 - GPU多卡训练:
+需要将 “/ppgan/modules/first_order.py”中的nn.BatchNorm 改为nn.SyncBatchNorm
 ```
 export CUDA_VISIBLE_DEVICES=0,1,2,3
 python -m paddle.distributed.launch \
@@ -133,7 +134,6 @@ python -m paddle.distributed.launch \
 
 
 ### 3. 模型压缩
-数据处理同上述，模型分为kp_detector和generator，首先固定原始generator部分，训练压缩版的kp_detector部分，然后固定原始kp_detector部分,去训练generator部分，最后将两个压缩的模型一起训练，同时添加中间的蒸馏loss。
 
 **预测:**
 ```
@@ -153,14 +153,19 @@ python -u tools/first-order-demo.py  \
 |       原始        |        229       |        0.012058867    |
 |       压缩        |        6.1       |      0.015025159    |
 
+**训练:**
+先将configs/firstorder_vox_mobile_256.yaml 中的mode设置成kp_detector, 训练压缩版
+的kp_detector的模型，固定原始generator模型；然后将configs/firstorder_vox_mobile_256.yaml 中的mode设置成generator，训练压缩版的generator的模型，固定原始kp_detector模型；最后将mode设置为both，修改配置文件中的kp_weight_path和gen_weight_path为>已经训练好的模型路径，一起训练。
+```
+export CUDA_VISIBLE_DEVICES=0
+python tools/main.py --config-file configs/firstorder_vox_mobile_256.yaml
+```
+
 ### 4. 模型部署
 #### 4.1 导出模型
 使用`tools/fom_export.py`脚本导出模型已经部署时使用的配置文件，配置文件名字为`firstorder_vox_mobile_256.yml`。模型导出脚本如下：
 ```bash
 # 导出FOM模型
-需要将 “/ppgan/modules/first_order.py”中的nn.SyncBatchNorm 改为nn.BatchNorm，因为export目前不支持SyncBatchNorm
-将 out = out[:, :, ::int_inv_scale, ::int_inv_scale] 改为 
-out = paddle.fluid.layers.resize_nearest(out, scale=self.scale)
 
 python tools/export_model.py \
     --config-file configs/firstorder_vox_mobile_256.yaml \
@@ -169,10 +174,10 @@ python tools/export_model.py \
     --export_model output_inference/
 ```
 预测模型会导出到`output_inference/fom_dy2st/`目录下，分别为`model.pdiparams`,  `model.pdiparams.info`, `model.pdmodel`。
-
+- [预训练模型](https://paddlegan.bj.bcebos.com/applications/first_order_model/paddle_lite/inference/lite.zip)
 
 #### 4.2 PaddleLite部署
-- [使用Paddle Lite部署FOM模型](./lite/README.md)
+- [使用Paddle Lite部署FOM模型](https://github.com/PaddlePaddle/PaddleGAN/tree/develop/deploy/lite)
 - [FOM-Lite-Demo](https://paddlegan.bj.bcebos.com/applications/first_order_model/paddle_lite/apk/face_detection_demo%202.zip)。更多内容，请参考[Paddle-Lite](https://github.com/PaddlePaddle/Paddle-Lite)
 目前问题：
 (a).Paddle Lite运行效果略差于Paddle Inference，正在优化中
@@ -190,4 +195,3 @@ python tools/export_model.py \
 }
 
 ```
-

ppgan/models/firstorder_model.py

@@ -31,6 +31,7 @@
 import cv2
 import os
 
+
 def init_weight(net):
     def reset_func(m):
         if isinstance(m, (nn.BatchNorm, nn.BatchNorm2D, nn.SyncBatchNorm)):
@@ -86,7 +87,7 @@ def setup_lr_schedulers(self, lr_cfg):
             "gen_lr": self.gen_lr,
             "dis_lr": self.dis_lr
         }
-    
+
     def setup_net_parallel(self):
         if isinstance(self.nets['Gen_Full'], paddle.DataParallel):
             self.nets['kp_detector'] = self.nets[
@@ -200,13 +201,14 @@ def forward(self, source, kp_source, kp_driving, kp_driving_initial):
                 kp_driving['jacobian'],
                 paddle.inverse(kp_driving_initial['jacobian']))
             kp_norm['jacobian'] = paddle.matmul(jacobian_diff,
-                                               kp_source['jacobian'])
-            out = self.generator(source, kp_source=kp_source, kp_driving=kp_norm)
+                                                kp_source['jacobian'])
+            out = self.generator(source,
+                                 kp_source=kp_source,
+                                 kp_driving=kp_norm)
             return out['prediction']
 
-
     def export_model(self, export_model=None, output_dir=None, inputs_size=[]):
-        
+
         source = paddle.rand(shape=inputs_size[0], dtype='float32')
         driving = paddle.rand(shape=inputs_size[1], dtype='float32')
         value = paddle.rand(shape=inputs_size[2], dtype='float32')
@@ -216,16 +218,172 @@ def export_model(self, export_model=None, output_dir=None, inputs_size=[]):
         driving1 = {'value': value, 'jacobian': j}
         driving2 = {'value': value2, 'jacobian': j2}
         driving3 = {'value': value, 'jacobian': j}
-        
+
         outpath = os.path.join(output_dir, "fom_dy2st")
         if not os.path.exists(outpath):
             os.makedirs(outpath)
-        paddle.jit.save(self.nets['Gen_Full'].kp_extractor, os.path.join(outpath, "kp_detector"), input_spec=[source])
+        paddle.jit.save(self.nets['Gen_Full'].kp_extractor,
+                        os.path.join(outpath, "kp_detector"),
+                        input_spec=[source])
         infer_generator = self.InferGenerator()
         infer_generator.set_generator(self.nets['Gen_Full'].generator)
-        paddle.jit.save(infer_generator, os.path.join(outpath, "generator"), input_spec=[source, driving1, driving2, driving3])
+        paddle.jit.save(infer_generator,
+                        os.path.join(outpath, "generator"),
+                        input_spec=[source, driving1, driving2, driving3])
+
+
+@MODELS.register()
+class FirstOrderModelMobile(FirstOrderModel):
+    """ This class implements the FirstOrderMotionMobile model, modified according to the FirstOrderMotion paper:
+    https://proceedings.neurips.cc/paper/2019/file/31c0b36aef265d9221af80872ceb62f9-Paper.pdf.
+    """
+    def __init__(self,
+                 common_params,
+                 train_params,
+                 generator_ori,
+                 generator,
+                 mode,
+                 kp_weight_path=None,
+                 gen_weight_path=None,
+                 discriminator=None):
+        super(FirstOrderModel, self).__init__()
+        modes = ["kp_detector", "generator", "both"]
+        assert mode in modes
+        # def local var
+        self.input_data = None
+        self.generated = None
+        self.losses_generator = None
+        self.train_params = train_params
 
+        # fix origin fom model for distill
+        generator_ori_cfg = generator_ori
+        generator_ori_cfg.update({'common_params': common_params})
+        generator_ori_cfg.update({'train_params': train_params})
+        generator_ori_cfg.update(
+            {'dis_scales': discriminator.discriminator_cfg.scales})
+        self.Gen_Full_ori = build_generator(generator_ori_cfg)
+        discriminator_cfg = discriminator
+        discriminator_cfg.update({'common_params': common_params})
+        discriminator_cfg.update({'train_params': train_params})
+        self.nets['Dis'] = build_discriminator(discriminator_cfg)
 
+        # define networks
+        generator_cfg = generator
+        generator_cfg.update({'common_params': common_params})
+        generator_cfg.update({'train_params': train_params})
+        generator_cfg.update(
+            {'dis_scales': discriminator.discriminator_cfg.scales})
+        if (mode == "kp_detector"):
+            print("just train kp_detector, fix generator")
+            generator_cfg.update(
+                {'generator_cfg': generator_ori_cfg['generator_cfg']})
+        elif mode == "generator":
+            print("just train generator, fix kp_detector")
+            generator_cfg.update(
+                {'kp_detector_cfg': generator_ori_cfg['kp_detector_cfg']})
+        elif mode == "both":
+            print("train both kp_detector and generator")
+        self.mode = mode
+        self.nets['Gen_Full'] = build_generator(generator_cfg)
+        self.kp_weight_path = kp_weight_path
+        self.gen_weight_path = gen_weight_path
+        self.visualizer = Visualizer()
+
+    def setup_net_parallel(self):
+        if isinstance(self.nets['Gen_Full'], paddle.DataParallel):
+            self.nets['kp_detector'] = self.nets[
+                'Gen_Full']._layers.kp_extractor
+            self.nets['generator'] = self.nets['Gen_Full']._layers.generator
+            self.kp_detector_ori = self.Gen_Full_ori._layers.kp_extractor
+            self.nets['generator'] = self.nets['Gen_Full']._layers.generator
+            self.nets['discriminator'] = self.nets['Dis']._layers.discriminator
+        else:
+            self.nets['kp_detector'] = self.nets['Gen_Full'].kp_extractor
+            self.nets['generator'] = self.nets['Gen_Full'].generator
+            self.kp_detector_ori = self.Gen_Full_ori.kp_extractor
+            self.nets['discriminator'] = self.nets['Dis'].discriminator
+
+        from ppgan.utils.download import get_path_from_url
+        vox_cpk_weight_url = 'https://paddlegan.bj.bcebos.com/applications/first_order_model/vox-cpk.pdparams'
+        weight_path = get_path_from_url(vox_cpk_weight_url)
+        checkpoint = paddle.load(weight_path)
+        if (self.mode == "kp_detector"):
+            self.nets['generator'].set_state_dict(checkpoint['generator'])
+            for param in self.nets['generator'].parameters():
+                param.stop_gradient = True
+        elif self.mode == "generator":
+            self.nets['kp_detector'].set_state_dict(checkpoint['kp_detector'])
+            for param in self.nets['kp_detector'].parameters():
+                param.stop_gradient = True
+        elif self.mode == "both":
+            checkpoint = paddle.load(self.kp_weight_path)
+            self.nets['kp_detector'].set_state_dict(checkpoint['kp_detector'])
+            checkpoint = paddle.load(self.gen_weight_path)
+            self.nets['generator'].set_state_dict(checkpoint['generator'])
+
+        self.kp_detector_ori.set_state_dict(checkpoint['kp_detector'])
+        for param in self.kp_detector_ori.parameters():
+            param.stop_gradient = True
+
+    def setup_optimizers(self, lr_cfg, optimizer):
+        self.setup_net_parallel()
+        # init params
+        init_weight(self.nets['discriminator'])
+        self.optimizers['optimizer_Dis'] = build_optimizer(
+            optimizer,
+            self.dis_lr,
+            parameters=self.nets['discriminator'].parameters())
+
+        if (self.mode == "kp_detector"):
+            init_weight(self.nets['kp_detector'])
+            self.optimizers['optimizer_KP'] = build_optimizer(
+                optimizer,
+                self.kp_lr,
+                parameters=self.nets['kp_detector'].parameters())
+        elif self.mode == "generator":
+            init_weight(self.nets['generator'])
+            self.optimizers['optimizer_Gen'] = build_optimizer(
+                optimizer,
+                self.gen_lr,
+                parameters=self.nets['generator'].parameters())
+        elif self.mode == "both":
+            super(FirstOrderModelMobile,
+                  self).setup_optimizers(lr_cfg, optimizer)
+        # define loss functions
+        self.losses = {}
+
+    def forward(self):
+        """Run forward pass; called by both functions <optimize_parameters> and <test>."""
+        if (self.mode == "kp_detector_distill"):
+            self.losses_generator, self.generated = \
+                self.nets['Gen_Full'](self.input_data.copy(), self.nets['discriminator'], self.kp_detector_ori)
+        else:
+            self.losses_generator, self.generated = \
+                self.nets['Gen_Full'](self.input_data.copy(), self.nets['discriminator'])
+
+    def train_iter(self, optimizers=None):
+        if (self.mode == "both"):
+            super(FirstOrderModelMobile, self).train_iter(optimizers=optimizers)
+            return
+        self.forward()
+        # update G
+        self.set_requires_grad(self.nets['discriminator'], False)
+        if (self.mode == "kp_detector"):
+            self.optimizers['optimizer_KP'].clear_grad()
+            self.backward_G()
+            self.optimizers['optimizer_KP'].step()
+        if (self.mode == "generator"):
+            self.optimizers['optimizer_Gen'].clear_grad()
+            self.backward_G()
+            self.optimizers['optimizer_Gen'].step()
+        outs = {}
+
+        # update D
+        if self.train_params['loss_weights']['generator_gan'] != 0:
+            self.set_requires_grad(self.nets['discriminator'], True)
+            self.optimizers['optimizer_Dis'].clear_grad()
+            self.backward_D()
+            self.optimizers['optimizer_Dis'].step()
 
 
 class Visualizer:

ppgan/models/generators/generator_firstorder.py

@@ -68,7 +68,7 @@ def __init__(self, generator_cfg, kp_detector_cfg, common_params,
         if sum(self.loss_weights['perceptual']) != 0:
             self.vgg = VGG19()
 
-    def forward(self, x, discriminator):
+    def forward(self, x, discriminator, kp_extractor_ori=None):
         kp_source = self.kp_extractor(x['source'])
         kp_driving = self.kp_extractor(x['driving'])
         generated = self.generator(x['source'],
@@ -151,6 +151,19 @@ def forward(self, x, discriminator):
                 value = paddle.abs(eye - value).mean()
                 loss_values['equivariance_jacobian'] = self.loss_weights[
                     'equivariance_jacobian'] * value
+
+        if kp_extractor_ori is not None:
+            recon_loss = paddle.nn.loss.L1Loss()
+
+            kp_distillation_loss_source = recon_loss(
+                kp_extractor_ori(x['source'])['value'],
+                self.kp_extractor(x['source'])['value'])
+            kp_distillation_loss_driving = recon_loss(
+                kp_extractor_ori(x['driving'])['value'],
+                self.kp_extractor(x['driving'])['value'])
+            loss_values[
+                "kp_distillation_loss"] = kp_distillation_loss_source + kp_distillation_loss_driving
+
         return loss_values, generated