add gan compression

demo/gan_compression/configs/__init__.py

+def encode_config(config):
+    return '_'.join([str(c) for c in config['channels']])
+
+
+def decode_config(config_str):
+    channels = config_str.split('_')
+    assert len(channels) == 6 or len(channels) == 8
+    channels = [int(c) for c in channels]
+    return {'channels': channels}

demo/gan_compression/configs/resnet_configs.py

+import random
+
+
+class ResnetConfigs:
+    def __init__(self, n_channels):
+        self.attributes = ['n_channels']
+        self.n_channels = n_channels
+
+    def sample(self):
+        ret = {}
+        ret['channels'] = []
+        for n_channel in self.n_channels:
+            ret['channels'].append(random.choice(n_channel))
+        return ret
+
+    def largest(self):
+        ret = {}
+        ret['channels'] = []
+        for n_channel in self.n_channels:
+            ret['channels'].append(max(n_channel))
+        return ret
+
+    def smallest(self):
+        ret = {}
+        ret['channels'] = []
+        for n_channel in self.n_channels:
+            ret['channels'].append(min(n_channel))
+        return ret
+
+    def all_configs(self):
+        def yield_channels(i):
+            if i == len(self.n_channels):
+                yield []
+                return
+            for n in self.n_channels[i]:
+                for after_channels in yield_channels(i + 1):
+                    yield [n] + after_channels
+
+        for channels in yield_channels(0):
+            yield {'channels': channels}
+
+    def __call__(self, name):
+        assert name in ('largest', 'smallest')
+        if name == 'largest':
+            return self.largest()
+        elif name == 'smallest':
+            return self.smallest()
+        else:
+            raise NotImplementedError
+
+    def __len__(self):
+        ret = 1
+        for n_channel in self.n_channels:
+            ret *= len(n_channel)
+
+
+def get_configs(config_name):
+    if config_name == 'channels-48':
+        return ResnetConfigs(
+            n_channels=[[48, 32], [48, 32], [48, 40, 32], [48, 40, 32],
+                        [48, 40, 32], [48, 40, 32], [48, 32, 24, 16],
+                        [48, 32, 24, 16]])
+    elif config_name == 'channels-32':
+        return ResnetConfigs(
+            n_channels=[[32, 24, 16], [32, 24, 16], [32, 24, 16],
+                        [32, 24, 16], [32, 24, 16], [32, 24, 16],
+                        [32, 24, 16], [32, 24, 16]])
+    elif config_name == 'debug':
+        return ResnetConfigs(n_channels=[[48, 32], [48, 32], [48, 40, 32],
+                                         [48, 40, 32], [48, 40, 32],
+                                         [48, 40, 32], [48], [48]])
+    elif config_name == 'test':
+        return ResnetConfigs(n_channels=[[8], [6, 8], [6, 8], [8], [8], [8],
+                                         [8], [8]])
+    else:
+        raise NotImplementedError('Unknown configuration [%s]!!!' %
+                                  config_name)

demo/gan_compression/configs/single_configs.py

+class SingleConfigs:
+    def __init__(self, config):
+        self.attributes = ['n_channels']
+        self.configs = [config]
+
+    def sample(self):
+        return self.configs[0]
+
+    def largest(self):
+        return self.configs[0]
+
+    def all_configs(self):
+        for config in self.configs:
+            yield config
+
+    def __call__(self, name):
+        assert name in ('largest', 'smallest')
+        if name == 'largest':
+            return self.largest()
+        elif name == 'smallest':
+            return self.smallest()
+        else:
+            raise NotImplementedError
+
+    def __str__(self):
+        ret = ''
+        for attr in self.attributes:
+            ret += 'attr: %s\n' % str(getattr(self, attr))
+        return ret
+
+    def __len__(self):
+        return len(self.configs)

demo/gan_compression/dataset/data_loader.py

+import paddle.fluid as fluid
+from data_reader import data_reader
+
+
+def create_data(cfgs, direction='AtoB', eval_mode=False):
+    if eval_mode == False:
+        mode = 'TRAIN'
+    else:
+        mode = 'EVAL'
+    reader = data_reader(cfgs, mode=mode)
+    data, id2name = reader.make_data(direction)
+    loader = fluid.io.DataLoader.from_generator(
+        capacity=4, iterable=True, use_double_buffer=True)
+
+    loader.set_batch_generator(
+        data,
+        places=fluid.CUDAPlace(0)
+        if cfgs.use_gpu else fluid.cpu_places())  ### fluid.cuda_places()
+    return loader, id2name
+
+
+def create_eval_data(cfgs, direction='AtoB'):
+    return create_data(cfgs, eval_mode=True)

demo/gan_compression/dataset/data_reader.py

+#copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+#Licensed under the Apache License, Version 2.0 (the "License");
+#you may not use this file except in compliance with the License.
+#You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+#Unless required by applicable law or agreed to in writing, software
+#distributed under the License is distributed on an "AS IS" BASIS,
+#WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#See the License for the specific language governing permissions and
+#limitations under the License.
+
+from __future__ import print_function
+from six.moves import range
+from PIL import Image, ImageOps
+
+import gzip
+import numpy as np
+import argparse
+import struct
+import os
+import paddle
+import paddle.fluid as fluid
+import random
+import sys
+
+
+def RandomCrop(img, crop_w, crop_h):
+    w, h = img.size[0], img.size[1]
+    i = np.random.randint(0, w - crop_w)
+    j = np.random.randint(0, h - crop_h)
+    return img.crop((i, j, i + crop_w, j + crop_h))
+
+
+def CentorCrop(img, crop_w, crop_h):
+    w, h = img.size[0], img.size[1]
+    i = int((w - crop_w) / 2.0)
+    j = int((h - crop_h) / 2.0)
+    return img.crop((i, j, i + crop_w, j + crop_h))
+
+
+def RandomHorizonFlip(img):
+    i = np.random.rand()
+    if i > 0.5:
+        img = ImageOps.mirror(img)
+    return img
+
+
+class reader_creator:
+    def __init__(self, *args, **kwcfgs):
+        raise NotImplementedError
+
+    def make_reader(self, cfgs):
+        raise NotImplementedError
+
+
+class single_datareader(reader_creator):
+    def __init__(self, list_filename, cfgs, mode='TEST'):
+        self.cfgs = cfgs
+        self.mode = mode
+        self.lines = open(list_filename).readlines()
+        (self.data_dir, _) = os.path.split(list_filename)
+        self.id2name = {}
+        if self.mode == "TRAIN":
+            self.shuffle = self.cfgs.shuffle
+        else:
+            self.shuffle = False
+
+    def len(self):
+        return self.len(self.lines)
+
+    def make_reader(self):
+        def reader():
+            batch_out_1 = []
+            batch_out_name = []
+            if self.shuffle:
+                np.random.shuffle(self.lines)
+
+            for i, fileA in enumerate(self.lines):
+                fileA = fileA.strip('\n\r\t').split(' ')[0]
+                self.id2name[i] = os.path.basename(fileA)
+                imgA = Image.open(os.path.join(self.data_dir, fileA)).convert(
+                    'RGB')
+
+                if self.mode == 'TRAIN':
+                    imgA = imgA.resize((self.cfgs.image_size,
+                                        self.cfgs.image_size), Image.BICUBIC)
+
+                    if self.cfgs.crop_type == 'Centor':
+                        imgA = CentorCrop(imgA, self.cfgs.crop_size,
+                                          self.cfgs.crop_size)
+                    elif self.cfgs.crop_type == 'Random':
+                        imgA = RandomCrop(imgA, self.cfgs.crop_size,
+                                          self.cfgs.crop_size)
+                    if self.cfgs.flip:
+                        imgA = RandomHorizonFlip(imgA)
+                else:
+                    imgA = imgA.resize((self.cfgs.crop_size,
+                                        self.cfgs.crop_size), Image.BICUBIC)
+
+                imgA = (np.array(imgA).astype('float32') / 255.0 - 0.5) / 0.5
+                imgA = imgA.transpose([2, 0, 1])
+
+                batch_out_1.append(imgA)
+                batch_out_name.append(i)
+
+                if len(batch_out_1) == self.cfgs.batch_size:
+                    yield batch_out_1, batch_out_name
+                    batch_out_1 = []
+                    batch_out_name = []
+
+        return reader
+
+
+class cycle_datareader(reader_creator):
+    def __init__(self, list_filename_A, list_filename_B, cfgs, mode='TRAIN'):
+        self.cfgs = cfgs
+        self.mode = mode
+        self.lines_A = open(list_filename_A).readlines()
+        self.lines_B = open(list_filename_B).readlines()
+        (self.data_dir, _) = os.path.split(list_filename_A)
+        self._max_dataset_size = max(len(self.lines_A), len(self.lines_B))
+        self.id2name = {}
+        if self.mode == "TRAIN":
+            if len(self.lines_A) < self._max_dataset_size:
+                rundant = self._max_dataset_size % len(self.lines_A)
+                self.lines_A.extend(self.lines_A[:rundant])
+            if len(self.lines_B) < self._max_dataset_size:
+                rundant = self._max_dataset_size % len(self.lines_B)
+                self.lines_B.extend(self.lines_B[:rundant])
+            self.shuffle = self.cfgs.shuffle
+        else:
+            self.shuffle = False
+
+    def len(self):
+        return self._max_dataset_size
+
+    def make_reader(self):
+        def reader():
+            batch_out_1 = []
+            batch_out_2 = []
+            batch_out_name = []
+
+            if self.shuffle:
+                np.random.shuffle(self.lines_B)
+
+            for i, (fileA,
+                    fileB) in enumerate(zip(self.lines_A, self.lines_B)):
+                fileA = fileA.strip('\n\r\t').split(' ')[0]
+                fileB = fileB.strip('\n\r\t').split(' ')[0]
+                self.id2name[i] = os.path.basename(
+                    fileA) + '###' + os.path.basename(fileB)
+
+                imgA = Image.open(os.path.join(self.data_dir, fileA)).convert(
+                    'RGB')
+                imgB = Image.open(os.path.join(self.data_dir, fileB)).convert(
+                    'RGB')
+
+                if self.mode == 'TRAIN':
+                    imgA = imgA.resize((self.cfgs.image_size,
+                                        self.cfgs.image_size), Image.BICUBIC)
+                    imgB = imgB.resize((self.cfgs.image_size,
+                                        self.cfgs.image_size), Image.BICUBIC)
+
+                    if self.cfgs.crop_type == 'Centor':
+                        imgA = CentorCrop(imgA, self.cfgs.crop_size,
+                                          self.cfgs.crop_size)
+                        imgB = CentorCrop(imgB, self.cfgs.crop_size,
+                                          self.cfgs.crop_size)
+                    elif self.cfgs.crop_type == 'Random':
+                        imgA = RandomCrop(imgA, self.cfgs.crop_size,
+                                          self.cfgs.crop_size)
+                        imgB = RandomCrop(imgB, self.cfgs.crop_size,
+                                          self.cfgs.crop_size)
+                    if self.cfgs.flip:
+                        imgA = RandomHorizonFlip(imgA)
+                        imgB = RandomHorizonFlip(imgB)
+                else:
+                    imgA = imgA.resize((self.cfgs.crop_size,
+                                        self.cfgs.crop_size), Image.BICUBIC)
+                    imgB = imgB.resize((self.cfgs.crop_size,
+                                        self.cfgs.crop_size), Image.BICUBIC)
+
+                imgA = (np.array(imgA).astype('float32') / 255.0 - 0.5) / 0.5
+                imgA = imgA.transpose([2, 0, 1])
+
+                imgB = (np.array(imgB).astype('float32') / 255.0 - 0.5) / 0.5
+                imgB = imgB.transpose([2, 0, 1])
+
+                batch_out_1.append(imgA)
+                batch_out_2.append(imgB)
+                batch_out_name.append(i)
+
+                if len(batch_out_1) == self.cfgs.batch_size:
+                    yield batch_out_1, batch_out_2, batch_out_name
+                    batch_out_2 = []
+                    batch_out_1 = []
+                    batch_out_name = []
+
+        return reader
+
+
+class data_reader(object):
+    def __init__(self, cfgs, mode='TRAIN'):
+        self.mode = mode
+        self.cfgs = cfgs
+
+    def make_data(self, direction='AtoB'):
+        if self.cfgs.model == 'cycle_gan':
+            dataset_dir = os.path.join(self.cfgs.dataroot, self.cfgs.dataset)
+            fileB_list = None
+            if self.mode == 'TRAIN':
+                fileA_list = os.path.join(dataset_dir, "trainA.txt")
+                fileB_list = os.path.join(dataset_dir, "trainB.txt")
+            else:
+                if direction == 'AtoB':
+                    fileA_list = os.path.join(dataset_dir, "testA.txt")
+                else:
+                    fileA_list = os.path.join(dataset_dir, "testB.txt")
+
+            if fileB_list is not None:
+                train_reader = cycle_datareader(
+                    list_filename_A=fileA_list,
+                    list_filename_B=fileB_list,
+                    cfgs=self.cfgs,
+                    mode=self.mode)
+            else:
+                train_reader = single_datareader(
+                    list_filename=fileA_list, cfgs=self.cfgs, mode=self.mode)
+
+            reader = train_reader.make_reader()
+            id2name = train_reader.id2name
+            return reader, id2name

demo/gan_compression/distillers/__init__.py

+import importlib
+from models.base_model import BaseModel
+
+
+def find_model_using_name(model_name):
+    model_filename = "distillers." + model_name + "_distiller"
+    modellib = importlib.import_module(model_filename)
+    target_model_name = model_name.replace("_", "") + "distiller"
+    model = None
+    for name, cls in modellib.__dict__.items():
+        if name.lower() == target_model_name.lower() and issubclass(cls,
+                                                                    BaseModel):
+            model = cls
+    assert model is not None, "model {} is not right, please check it!".format(
+        model_name)
+
+    return model
+
+
+def get_special_cfg(model):
+    model_cls = find_model_using_name(model)
+    return model_cls.add_special_cfgs
+
+
+def create_distiller(cfg):
+    distiller = find_model_using_name(cfg.distiller)
+    return distiller(cfg)

demo/gan_compression/distillers/base_resnet_distiller.py

+import os
+import numpy as np
+import itertools
+import paddle.fluid as fluid
+from paddle.fluid.dygraph.nn import Conv2D
+from paddle.fluid.dygraph.base import to_variable
+from models.super_modules import SuperConv2D
+from models import loss
+from models import network
+from models.base_model import BaseModel
+from utils import util, optimization
+from dataset.data_loader import create_eval_data
+from metric.inception import InceptionV3
+
+
+class BaseResnetDistiller(BaseModel):
+    @staticmethod
+    def add_special_cfgs(parser, **kwcfgs):
+        parser.add_argument(
+            '--lambda_distill',
+            type=float,
+            default=0.01,
+            help="the initialize lambda parameter for distiller loss")
+        parser.add_argument(
+            '--lambda_gan',
+            type=float,
+            default=1,
+            help="the initialize lambda parameter for gan loss")
+        parser.add_argument(
+            '--lambda_recon',
+            type=float,
+            default=10.0,
+            help="the initialize lambda parameter for recon loss")
+        parser.add_argument(
+            '--student_ngf',
+            type=int,
+            default=32,
+            help="base channel of student generator")
+        parser.add_argument(
+            '--student_dropout_rate',
+            type=float,
+            default=0,
+            help="dropout rate of student generator")
+        parser.add_argument(
+            '--restore_teacher_G_path',
+            type=str,
+            default=None,
+            help="the pretrain model path of teacher generator")
+        parser.add_argument(
+            '--restore_student_G_path',
+            type=str,
+            default=None,
+            help="the pretrain model path of student generator")
+        parser.add_argument(
+            '--real_stat_path',
+            type=str,
+            default='real_stat/horse2zebra_B.npz',
+            help="path of real stat")
+        return parser
+
+    def __init__(self, cfgs, task):
+        super(BaseResnetDistiller, self).__init__()
+        self.cfgs = cfgs
+        self.task = task
+        self.loss_names = ['G_gan', 'G_distill', 'G_recon', 'D_fake', 'D_real']
+        self.model_names = ['netG_student', 'netG_teacher', 'netD']
+        self.netG_teacher = network.define_G(cfgs.input_nc, cfgs.output_nc,
+                                             cfgs.ngf, cfgs.netG,
+                                             cfgs.norm_type, cfgs.dropout_rate)
+        student_netG = getattr(cfgs, self.task + '_student_netG')
+        self.netG_student = network.define_G(
+            cfgs.input_nc, cfgs.output_nc, cfgs.student_ngf, student_netG,
+            cfgs.norm_type, cfgs.student_dropout_rate)
+        if self.task == 'distiller':
+            self.netG_pretrained = network.define_G(
+                cfgs.input_nc, cfgs.output_nc, cfgs.pretrained_ngf,
+                cfgs.pretrained_netG, cfgs.norm_type, 0)
+
+        self.netD = network.define_D(cfgs.output_nc, cfgs.ndf, cfgs.netD,
+                                     cfgs.norm_type, cfgs.n_layer_D)
+
+        self.netG_teacher.eval()
+        self.netG_student.train()
+        self.netD.train()
+
+        ### [9, 12, 15, 18]
+        self.mapping_layers = ['model.%d' % i for i in range(9, 21, 3)]
+        self.netAs = []
+        self.Tacts, self.Sacts = {}, {}
+
+        G_params = self.netG_student.parameters()
+        for i, n in enumerate(self.mapping_layers):
+            ft, fs = cfgs.ngf, cfgs.student_ngf
+            if self.task == 'distiller':
+                netA = Conv2D(
+                    num_channels=fs * 4, num_filters=ft * 4, filter_size=1)
+            else:
+                netA = SuperConv2D(
+                    num_channels=fs * 4, num_filters=ft * 4, filter_size=1)
+
+            G_params += netA.parameters()
+            self.netAs.append(netA)
+            self.loss_names.append('G_distill%d' % i)
+
+        if self.task == 'distiller':
+            learning_rate = cfgs.distiller_lr
+            scheduler = cfgs.distiller_scheduler
+            nepochs = cfgs.distiller_nepochs
+            nepochs_decay = cfgs.distiller_nepochs_decay
+        elif self.task == 'supernet':
+            learning_rate = cfgs.supernet_lr
+            scheduler = cfgs.supernet_scheduler
+            nepochs = cfgs.supernet_nepochs
+            nepochs_decay = cfgs.supernet_nepochs_decay
+        else:
+            raise NotImplementedError("task {} is not suppport".format(
+                self.task))
+
+        self.optimizer_G = optimization.Optimizer(
+            learning_rate,
+            scheduler,
+            cfgs.step_per_epoch,
+            nepochs,
+            nepochs_decay,
+            cfgs,
+            parameter_list=G_params)
+        self.optimizer_D = optimization.Optimizer(
+            learning_rate,
+            scheduler,
+            cfgs.step_per_epoch,
+            nepochs,
+            nepochs_decay,
+            cfgs,
+            parameter_list=self.netD.parameters())
+        self.eval_dataloader, self.name = create_eval_data(
+            cfgs, direction=cfgs.direction)
+
+        block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[2048]
+        self.inception_model = InceptionV3([block_idx])
+        self.is_best = False
+
+        if cfgs.real_stat_path:
+            self.npz = np.load(cfgs.real_stat_path)
+
+        self.is_best = False
+
+    def setup(self):
+        self.load_networks()
+
+        if self.cfgs.lambda_distill > 0:
+
+            def get_activation(mem, name):
+                def get_output_hook(layer, input, output):
+                    mem[name] = output
+
+                return get_output_hook
+
+            def add_hook(net, mem, mapping_layers):
+                for idx, (n, m) in enumerate(net.named_sublayers()):
+                    if n in mapping_layers:
+                        m.register_forward_post_hook(get_activation(mem, n))
+
+        add_hook(self.netG_teacher, self.Tacts, self.mapping_layers)
+        add_hook(self.netG_student, self.Sacts, self.mapping_layers)
+
+    def set_input(self, inputs):
+        self.real_A = inputs[0] if self.cfgs.direction == 'AtoB' else inputs[1]
+        self.real_B = inputs[1] if self.cfgs.direction == 'AtoB' else inputs[0]
+
+    def set_single_input(self, inputs):
+        self.real_A = inputs[0]
+
+    def load_networks(self):
+        if self.cfgs.restore_teacher_G_path is None:
+            if self.cfgs.direction == 'AtoB':
+                teacher_G_path = os.path.join(self.cfgs.save_dir, 'mobile',
+                                              'last_netG_A')
+            else:
+                teacher_G_path = os.path.join(self.cfgs.save_dir, 'mobile',
+                                              'last_netG_B')
+        else:
+            teacher_G_path = self.cfgs.restore_teacher_G_path
+
+        util.load_network(self.netG_teacher, teacher_G_path)
+
+        if self.cfgs.restore_student_G_path is not None:
+            util.load_network(self.netG_student,
+                              self.cfgs.restore_student_G_path)
+        else:
+            if self.task == 'supernet':
+                student_G_path = os.path.join(self.cfgs.save_dir, 'distiller',
+                                              'last_stu_netG')
+                util.load_network(self.netG_student, student_G_path)
+
+        if self.cfgs.restore_D_path is not None:
+            util.load_network(self.netD, self.cfgs.restore_D_path)
+        if self.cfgs.restore_A_path is not None:
+            for i, netA in enumerate(self.netAs):
+                netA_path = '%s-%d.pth' % (self.cfgs.restore_A_path, i)
+                util.load_network(netA, netA_path)
+        if self.cfgs.restore_O_path is not None:
+            util.load_optimizer(self.optimizer_G,
+                                self.cfgs.restore_G_optimizer_path)
+            util.load_optimizer(self.optimizer_D,
+                                self.cfgs.restore_D_optimizer_path)
+
+    def backward_D(self):
+        fake = self.Sfake_B.detach()
+        real = self.real_B.detach()
+
+        pred_fake = self.netD(fake)
+        self.loss_D_fake = loss.gan_loss(self.cfgs.gan_loss_mode, pred_fake,
+                                         False)
+
+        pred_real = self.netD(real)
+        self.loss_D_real = loss.gan_loss(self.cfgs.gan_loss_mode, pred_real,
+                                         True)
+
+        self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5
+        self.loss_D.backward()
+
+    def calc_distill_loss(self):
+        raise NotImplementedError
+
+    def backward_G(self):
+        raise NotImplementedError
+
+    def optimize_parameter(self):
+        raise NotImplementedError
+
+    def forward(self):
+        raise NotImplementedError
+
+    def set_stop_gradient(self, nets, stop_grad=False):
+        if not isinstance(nets, list):
+            nets = [nets]
+        for net in nets:
+            if net is not None:
+                for param in net.parameters():
+                    param.stop_gradient = stop_grad
+
+    def save_network(self, epoch):
+        save_filename = '{}_stu_netG'.format(epoch)
+        save_path = os.path.join(self.cfgs.save_dir, self.task, save_filename)
+        fluid.save_dygraph(self.netG_student.state_dict(), save_path)
+        save_filename = '{}_tea_netG'.format(epoch)
+        save_path = os.path.join(self.cfgs.save_dir, self.task, save_filename)
+        fluid.save_dygraph(self.netG_teacher.state_dict(), save_path)
+
+        save_filename = '{}_netD'.format(epoch)
+        save_path = os.path.join(self.cfgs.save_dir, self.task, save_filename)
+        fluid.save_dygraph(self.netD.state_dict(), save_path)
+
+        for idx, net in enumerate(self.netAs):
+            save_filename = '{}_netA-{}'.format(epoch, idx)
+            save_path = os.path.join(self.cfgs.save_dir, self.task,
+                                     save_filename)
+            fluid.save_dygraph(net.state_dict(), save_path)
+
+    def get_current_loss(self):
+        loss_dict = {}
+        for name in self.loss_names:
+            if not hasattr(self, 'loss_' + name):
+                continue
+            key = name
+
+            def has_num(key):
+                for i in range(10):
+                    if str(i) in key:
+                        return True
+                return False
+
+            if has_num(key):
+                continue
+            loss_dict[key] = float(getattr(self, 'loss_' + name))
+        return loss_dict
+
+    def get_current_lr(self):
+        lr_dict = {}
+        lr_dict['optim_G'] = self.optimizer_G.optimizer.current_step_lr()
+        lr_dict['optim_D'] = self.optimizer_D.optimizer.current_step_lr()
+        return lr_dict
+
+    def test(self):
+        with fluid.dygraph.no_grad():
+            self.forward()

demo/gan_compression/distillers/resnet_distiller.py

+import os
+import numpy as np
+import paddle.fluid as fluid
+from paddle.fluid.dygraph.nn import Conv2D
+from .base_resnet_distiller import BaseResnetDistiller
+from utils import util
+from utils.weight_transfer import load_pretrained_weight
+from metric import compute_fid
+from models import loss
+from metric import get_fid
+
+
+class ResnetDistiller(BaseResnetDistiller):
+    @staticmethod
+    def add_special_cfgs(parser, load_pre=False):
+        parser.add_argument(
+            '--distiller_lr',
+            type=float,
+            default=2e-4,
+            help="Initial learning rate in train distiller net")
+        parser.add_argument(
+            '--distiller_epoch',
+            type=int,
+            default=200,
+            help="The number of epoch to train distiller net")
+        parser.add_argument(
+            '--distiller_nepochs',
+            type=int,
+            default=100,
+            help="number of epochs with the initial learning rate")
+        parser.add_argument(
+            '--distiller_nepochs_decay',
+            type=int,
+            default=100,
+            help="number of epochs to linearly decay learning rate to zero")
+        parser.add_argument(
+            '--distiller_scheduler',
+            type=str,
+            default='linear',
+            help="learning rate scheduler in train distiller net")
+        parser.add_argument(
+            '--distiller_student_netG',
+            type=str,
+            default='mobile_resnet_9blocks',
+            help="Which student generator network to choose in distiller")
+        parser.add_argument(
+            '--pretrained_ngf',
+            type=int,
+            default=64,
+            help="Base channels in generator")
+        parser.add_argument(
+            '--pretrained_netG',
+            type=str,
+            default='mobile_resnet_9blocks',
+            help="Which generator network to choose in pretrain model")
+        parser.add_argument(
+            '--restore_pretrained_G_path',
+            type=str,
+            default=None,
+            help="the pretrain model of pretrain_model used in distiller")
+        if load_pre:
+            super(ResnetDistiller, ResnetDistiller).add_special_cfgs(parser)
+        return parser
+
+    def __init__(self, cfgs):
+        super(ResnetDistiller, self).__init__(cfgs, task='distiller')
+        self.best_fid = 1e9
+        self.fids = []
+        self.npz = np.load(cfgs.real_stat_path)
+
+    def forward(self):
+        with fluid.dygraph.no_grad():
+            self.Tfake_B = self.netG_teacher(self.real_A)
+        self.Sfake_B = self.netG_student(self.real_A)
+
+    def calc_distill_loss(self):
+        losses = []
+        for i, netA in enumerate(self.netAs):
+            assert isinstance(netA, Conv2D)
+            n = self.mapping_layers[i]
+            Tact = self.Tacts[n]
+            Tact.stop_gradient = True
+            Sact = self.Sacts[n]
+            ### 1x1 conv to match channels
+            Sact = netA(Sact)
+            loss = fluid.layers.mse_loss(Sact, Tact)
+            setattr(self, 'loss_G_distill%d' % i, loss)
+            losses.append(loss)
+        return sum(losses)
+
+    def backward_G(self):
+        self.loss_G_recon = loss.recon_loss(
+            self.cfgs.recon_loss_mode, self.Sfake_B,
+            self.Tfake_B) * self.cfgs.lambda_recon
+        pred_fake = self.netD(self.Sfake_B)
+        self.loss_G_gan = loss.gan_loss(
+            self.cfgs.gan_loss_mode, pred_fake, True,
+            for_discriminator=False) * self.cfgs.lambda_gan
+        if self.cfgs.lambda_distill > 0:
+            self.loss_G_distill = self.calc_distill_loss(
+            ) * self.cfgs.lambda_distill
+        else:
+            self.loss_G_distill = 0
+
+        self.loss_G = self.loss_G_gan + self.loss_G_recon + self.loss_G_distill
+        self.loss_G.backward()
+
+    def optimize_parameter(self):
+        self.forward()
+
+        self.set_stop_gradient(self.netD, False)
+        self.backward_D()
+
+        self.set_stop_gradient(self.netD, True)
+        self.backward_G()
+        self.optimizer_D.optimizer.minimize(self.loss_D)
+        self.optimizer_D.optimizer.clear_gradients()
+        self.optimizer_G.optimizer.minimize(self.loss_G)
+        self.optimizer_G.optimizer.clear_gradients()
+
+    def load_networks(self):
+        load_pretrain = True
+        if self.cfgs.restore_pretrained_G_path is not None:
+            pretrained_G_path = self.cfgs.restore_pretrained_G_path
+        else:
+            pretrained_G_path = os.path.join(self.cfgs.save_dir, 'mobile',
+                                             'last_netG_B')
+            if not os.path.exists(os.path.join(pretrained_G_path, 'pdparams')):
+                load_pretrain = False
+
+        if load_pretrain:
+            util.load_network(self.netG_pretrained, pretrained_G_path)
+            load_pretrained_weight(
+                self.cfgs.pretrained_netG, self.cfgs.student_netG,
+                self.netG_pretrained, self.netG_student,
+                self.cfgs.pretrained_ngf, self.cfgs.student_ngf)
+            del self.netG_pretrained
+        super(ResnetDistiller, self).load_networks()
+
+    def evaluate_model(self, step):
+        ret = {}
+        self.is_best = False
+        save_dir = os.path.join(self.cfgs.save_dir, 'distiller', 'eval',
+                                str(step))
+        if not os.path.exists(save_dir):
+            os.makedirs(save_dir)
+        self.netG_student.eval()
+        fakes = []
+        cnt = 0
+        for i, data_i in enumerate(self.eval_dataloader):
+            id2name = self.name
+            self.set_single_input(data_i)
+            self.test()
+            fakes.append(self.Sfake_B.detach().numpy())
+            for j in range(len(self.Sfake_B)):
+                if cnt < 10:
+                    Sname = 'Sfake_' + str(id2name[i + j]) + '.png'
+                    Tname = 'Tfake_' + str(id2name[i + j]) + '.png'
+                    Sfake_im = util.tensor2img(self.Sfake_B[j])
+                    Tfake_im = util.tensor2img(self.Tfake_B[j])
+                    util.save_image(Sfake_im, os.path.join(save_dir, Sname))
+                    util.save_image(Tfake_im, os.path.join(save_dir, Tname))
+                cnt += 1
+
+        suffix = self.cfgs.direction
+        fluid.disable_imperative()
+        fid = get_fid(fakes, self.inception_model, self.npz,
+                      self.cfgs.inception_model)
+        fluid.enable_imperative()
+        if fid < self.best_fid:
+            self.is_best = True
+            self.best_fid = fid
+        print("fid score is: %f, best fid score is %f" % (fid, self.best_fid))
+        self.fids.append(fid)
+        if len(self.fids) > 3:
+            self.fids.pop(0)
+        ret['metric/fid'] = fid
+        ret['metric/fid-mean'] = sum(self.fids) / len(self.fids)
+        ret['metric/fid-best'] = self.best_fid
+
+        self.netG_student.train()
+        return ret

demo/gan_compression/gan_compression.py

+import os
+import time
+import logging
+
+import paddle.fluid as fluid
+from dataset.data_loader import create_data
+from utils.get_args import configs
+
+
+class gan_compression:
+    def __init__(self, cfgs, **kwargs):
+        self.cfgs = cfgs
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+    def start_train(self):
+        steps = self.cfgs.task.split('+')
+        for step in steps:
+            if step == 'mobile':
+                from models import create_model
+            elif step == 'distiller':
+                from distiller import create_distiller as create_model
+            elif step == 'supernet':
+                from supernet import create_supernet as create_model
+            else:
+                raise NotImplementedError
+
+            print(
+                "============================= start train {} ==============================".
+                format(step))
+            fluid.enable_imperative()
+            model = create_model(self.cfgs)
+            model.setup()
+
+            _train_dataloader, _ = create_data(self.cfgs)
+
+            epochs = getattr(self.cfgs, '{}_epoch'.format(step))
+
+            for epoch_id in range(epochs):
+                for batch_id, data in enumerate(_train_dataloader()):
+                    start_time = time.time()
+                    model.set_input(data)
+                    model.optimize_parameter()
+                    batch_time = time.time() - start_time
+                    if batch_id % self.cfgs.print_freq == 0:
+                        message = 'epoch: %d, batch: %d batch_time: %fs' % (
+                            epoch_id, batch_id, batch_time)
+                        for k, v in model.get_current_lr().items():
+                            message += '%s: %f ' % (k, v)
+                        message += '\n'
+                        for k, v in model.get_current_loss().items():
+                            message += '%s: %.3f ' % (k, v)
+                        logging.info(message)
+
+                if epoch_id % self.cfgs.save_freq == 0 or epoch_id == (
+                        epochs - 1):
+                    model.evaluate_model(epoch_id)
+                    model.save_network(epoch_id)
+                    if epoch_id == (epochs - 1):
+                        model.save_network('last')
+            print("=" * 80)
+
+
+if __name__ == '__main__':
+    cfg_instance = configs()
+    cfgs = cfg_instance.get_all_config()
+    cfg_instance.print_configs(cfgs)
+    compression = gan_compression(cfgs)
+    compression.start_train()

demo/gan_compression/metric/__init__.py

+import numpy as np
+from metric.compute_fid import _compute_statistic_of_img, _calculate_frechet_distance
+from utils import util
+
+
+def get_fid(fakes, model, npz, premodel_path, batch_size=1, use_gpu=True):
+    m1, s1 = npz['mu'], npz['sigma']
+    fakes = np.concatenate(fakes, axis=0)
+    fakes = util.tensor2img(fakes).astype('float32')
+    m2, s2 = _compute_statistic_of_img(fakes, model, batch_size, 2048, use_gpu,
+                                       premodel_path)
+    return float(_calculate_frechet_distance(m1, s1, m2, s2))

demo/gan_compression/metric/compute_fid.py

+import numpy as np
+import os
+import fnmatch
+import cv2
+from cv2 import imread
+import paddle.fluid as fluid
+from scipy import linalg
+from inception import InceptionV3
+import requests
+
+
+def tqdm(x):
+    return x
+
+
+""" based on https://github.com/mit-han-lab/gan-compression/blob/master/metric/fid_score.py
+"""
+"""
+inceptionV3 pretrain model is convert from pytorch, pretrain_model url is 
+"""
+
+
+def _calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
+    m1 = np.atleast_1d(mu1)
+    m2 = np.atleast_1d(mu2)
+
+    sigma1 = np.atleast_2d(sigma1)
+    sigma2 = np.atleast_2d(sigma2)
+
+    assert mu1.shape == mu2.shape, 'Training and test mean vectors have different lengths'
+    assert sigma1.shape == sigma2.shape, 'Training and test covariances have different dimensions'
+
+    diff = mu1 - mu2
+
+    t = sigma1.dot(sigma2)
+    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
+    if not np.isfinite(covmean).all():
+        msg = ('fid calculation produces singular product; '
+               'adding %s to diagonal of cov estimates') % eps
+        print(msg)
+        offset = np.eye(sigma1.shape[0]) * eps
+        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
+
+    # Numerical error might give slight imaginary component
+    if np.iscomplexobj(covmean):
+        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
+            m = np.max(np.abs(covmean.imag))
+            raise ValueError('Imaginary component {}'.format(m))
+        covmean = covmean.real
+
+    tr_covmean = np.trace(covmean)
+
+    return (
+        diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean)
+
+
+def _build_program(model):
+    main_program = fluid.Program()
+    startup_program = fluid.Program()
+    with fluid.program_guard(main_program, startup_program):
+        images = fluid.data(name='images', shape=[None, 3, None, None])
+        output = model.network(images, class_dim=1008)
+        pred = fluid.layers.pool2d(output[0], global_pooling=True)
+
+    test_program = main_program.clone(for_test=True)
+    return pred, test_program, startup_program
+
+
+def _get_activations_from_ims(img, model, batch_size, dims, use_gpu,
+                              premodel_path):
+    n_batches = (len(img) + batch_size - 1) // batch_size
+    n_used_img = len(img)
+
+    pred_arr = np.empty((n_used_img, dims))
+
+    for i in tqdm(range(n_batches)):
+        start = i * batch_size
+        end = start + batch_size
+        if end > len(img):
+            end = len(img)
+        images = img[start:end]
+        if images.shape[1] != 3:
+            images = images.transpose((0, 3, 1, 2))
+        images /= 255
+
+        output, main_program, startup_program = _build_program(model)
+        place = fluid.CUDAPlace(0) if use_gpu else fluid.CPUPlace()
+        exe = fluid.Executor(place)
+        exe.run(startup_program)
+
+        fluid.load(main_program,
+                   os.path.join(premodel_path, 'paddle_inceptionv3'), exe)
+        pred = exe.run(main_program,
+                       feed={'images': images},
+                       fetch_list=[output])[0]
+
+        pred_arr[start:end] = pred.reshape(end - start, -1)
+
+    return pred_arr
+
+
+def _compute_statistic_of_img(img, model, batch_size, dims, use_gpu,
+                              premodel_path):
+    act = _get_activations_from_ims(img, model, batch_size, dims, use_gpu,
+                                    premodel_path)
+    mu = np.mean(act, axis=0)
+    sigma = np.cov(act, rowvar=False)
+    return mu, sigma
+
+
+def calculate_fid_given_img(img_fake,
+                            img_real,
+                            batch_size,
+                            use_gpu,
+                            dims,
+                            premodel_path,
+                            model=None):
+    assert os.path.exists(
+        premodel_path
+    ), 'pretrain_model path {} is not exists! Please download it first'.format(
+        premodel_path)
+    if model is None:
+        block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]
+        model = InceptionV3([block_idx])
+
+    m1, s1 = _compute_statistic_of_img(img_fake, model, batch_size, dims,
+                                       use_gpu, premodel_path)
+    m2, s2 = _compute_statistic_of_img(img_real, model, batch_size, dims,
+                                       use_gpu, premodel_path)
+
+    fid_value = _calculate_frechet_distance(m1, s1, m2, s2)
+    return fid_value
+
+
+def _get_activations(files, model, batch_size, dims, use_gpu, premodel_path):
+    if len(files) % batch_size != 0:
+        print(('Warning: number of images is not a multiple of the '
+               'batch size. Some samples are going to be ignored.'))
+    if batch_size > len(files):
+        print(('Warning: batch size is bigger than the datasets size. '
+               'Setting batch size to datasets size'))
+        batch_size = len(files)
+
+    n_batches = len(files) // batch_size
+    n_used_imgs = n_batches * batch_size
+
+    pred_arr = np.empty((n_used_imgs, dims))
+    for i in tqdm(range(n_batches)):
+        start = i * batch_size
+        end = start + batch_size
+        images = np.array(
+            [imread(str(f)).astype(np.float32) for f in files[start:end]])
+
+        if len(images.shape) != 4:
+            images = imread(str(files[start]))
+            images = cv2.cvtColor(images, cv2.COLOR_BGR2GRAY)
+            images = np.array([images.astype(np.float32)])
+
+        images = images.transpose((0, 3, 1, 2))
+        images /= 255
+
+        output, main_program, startup_program = _build_program(model)
+        place = fluid.CUDAPlace(0) if use_gpu else fluid.CPUPlace()
+        exe = fluid.Executor(place)
+        exe.run(startup_program)
+
+        fluid.load(main_program,
+                   os.path.join(premodel_path, 'paddle_inceptionv3'), exe)
+        pred = exe.run(main_program,
+                       feed={'images': images},
+                       fetch_list=[output])[0]
+
+        pred_arr[start:end] = pred.reshape(end - start, -1)
+
+    return pred_arr
+
+
+def _calculate_activation_statistics(files,
+                                     model,
+                                     premodel_path,
+                                     batch_size=50,
+                                     dims=2048,
+                                     use_gpu=False):
+    act = _get_activations(files, model, batch_size, dims, use_gpu,
+                           premodel_path)
+    mu = np.mean(act, axis=0)
+    sigma = np.cov(act, rowvar=False)
+    return mu, sigma
+
+
+def _compute_statistics_of_path(path, model, batch_size, dims, use_gpu,
+                                premodel_path):
+    if path.endswith('.npz'):
+        f = np.load(path)
+        m, s = f['mu'][:], f['sigma'][:]
+        f.close()
+    else:
+        files = []
+        for root, dirnames, filenames in os.walk(path):
+            for filename in fnmatch.filter(
+                    filenames, '*.jpg') or fnmatch.filter(filenames, '*.png'):
+                files.append(os.path.join(root, filename))
+        m, s = _calculate_activation_statistics(files, model, premodel_path,
+                                                batch_size, dims, use_gpu)
+    return m, s
+
+
+def calculate_fid_given_paths(paths,
+                              batch_size,
+                              use_gpu,
+                              dims,
+                              premodel_path,
+                              model=None):
+    assert os.path.exists(
+        premodel_path
+    ), 'pretrain_model path {} is not exists! Please download it first'.format(
+        premodel_path)
+    for p in paths:
+        if not os.path.exists(p):
+            raise RuntimeError('Invalid path: %s' % p)
+
+    if model is None:
+        block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]
+        model = InceptionV3([block_idx])
+
+    m1, s1 = _compute_statistics_of_path(paths[0], model, batch_size, dims,
+                                         use_gpu, premodel_path)
+    m2, s2 = _compute_statistics_of_path(paths[1], model, batch_size, dims,
+                                         use_gpu, premodel_path)
+
+    fid_value = _calculate_frechet_distance(m1, s1, m2, s2)
+    return fid_value

demo/gan_compression/metric/inception.py

+#copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+#Licensed under the Apache License, Version 2.0 (the "License");
+#you may not use this file except in compliance with the License.
+#You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+#Unless required by applicable law or agreed to in writing, software
+#distributed under the License is distributed on an "AS IS" BASIS,
+#WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#See the License for the specific language governing permissions and
+#limitations under the License.
+
+import math
+import paddle
+import paddle.fluid as fluid
+from paddle.fluid.param_attr import ParamAttr
+
+__all__ = ['InceptionV3']
+
+
+class InceptionV3:
+    DEFAULT_BLOCK_INDEX = 3
+    BLOCK_INDEX_BY_DIM = {
+        64: 0,  # First max pooling features
+        192: 1,  # Second max pooling featurs
+        768: 2,  # Pre-aux classifier features
+        2048: 3  # Final average pooling features
+    }
+
+    def __init__(self,
+                 output_blocks=[DEFAULT_BLOCK_INDEX],
+                 resize_input=True,
+                 normalize_input=True):
+        self.resize_input = resize_input
+        self.normalize_input = normalize_input
+        self.output_blocks = sorted(output_blocks)
+        self.last_needed_block = max(output_blocks)
+
+        assert self.last_needed_block <= 3, 'Last possible output block index is 3'
+
+    def network(self, x, class_dim=1000, aux_logits=False):
+        if self.resize_input:
+            x = fluid.layers.resize_bilinear(
+                x, out_shape=[299, 299], align_corners=False, align_mode=0)
+
+        if self.normalize_input:
+            x = x * 2 - 1
+
+        out, _, = self.fid_inceptionv3(x, class_dim, aux_logits)
+        return out
+
+    def fid_inceptionv3(self, x, num_classes=1000, aux_logits=False):
+        """ inception v3 model for FID computation
+        """
+        out = []
+        aux = None
+
+        ### block0
+        x = self.conv_bn_layer(x, 32, 3, stride=2, name='Conv2d_1a_3x3')
+        x = self.conv_bn_layer(x, 32, 3, name='Conv2d_2a_3x3')
+        x = self.conv_bn_layer(x, 64, 3, padding=1, name='Conv2d_2b_3x3')
+        x = fluid.layers.pool2d(x, pool_size=3, pool_stride=2, pool_type='max')
+        if 0 in self.output_blocks:
+            out.append(x)
+
+        if self.last_needed_block >= 1:
+            ### block1
+            x = self.conv_bn_layer(x, 80, 1, name='Conv2d_3b_1x1')
+            x = self.conv_bn_layer(x, 192, 3, name='Conv2d_4a_3x3')
+            x = fluid.layers.pool2d(
+                x, pool_size=3, pool_stride=2, pool_type='max')
+            if 1 in self.output_blocks:
+                out.append(x)
+
+        if self.last_needed_block >= 2:
+            ### block2
+            ### Mixed_5b 5c 5d
+            x = self.fid_inceptionA(x, pool_features=32, name='Mixed_5b')
+            x = self.fid_inceptionA(x, pool_features=64, name='Mixed_5c')
+            x = self.fid_inceptionA(x, pool_features=64, name='Mixed_5d')
+
+            ### Mixed_6
+            x = self.inceptionB(x, name='Mixed_6a')
+            x = self.fid_inceptionC(x, c7=128, name='Mixed_6b')
+            x = self.fid_inceptionC(x, c7=160, name='Mixed_6c')
+            x = self.fid_inceptionC(x, c7=160, name='Mixed_6d')
+            x = self.fid_inceptionC(x, c7=192, name='Mixed_6e')
+            if 2 in self.output_blocks:
+                out.append(x)
+
+            if aux_logits:
+                aux = self.inceptionAux(x, num_classes, name='AuxLogits')
+
+        if self.last_needed_block >= 3:
+            ### block3
+            ### Mixed_7
+            x = self.inceptionD(x, name='Mixed_7a')
+            x = self.fid_inceptionE_1(x, name='Mixed_7b')
+            x = self.fid_inceptionE_2(x, name='Mixed_7c')
+
+            x = fluid.layers.pool2d(x, global_pooling=True, pool_type='avg')
+            out.append(x)
+
+            #x = fluid.layers.dropout(x, dropout_prob=0.5)
+            #x = fluid.layers.flatten(x, axis=1)
+            #x = fluid.layers.fc(x, size=num_classes, param_attr=ParamAttr(name='fc.weight'), bias_attr=ParamAttr(name='fc.bias'))
+
+        return out, aux
+
+    def inceptionA(self, x, pool_features, name=None):
+        branch1x1 = self.conv_bn_layer(x, 64, 1, name=name + '.branch1x1')
+
+        branch5x5 = self.conv_bn_layer(x, 48, 1, name=name + '.branch5x5_1')
+        branch5x5 = self.conv_bn_layer(
+            branch5x5, 64, 5, padding=2, name=name + '.branch5x5_2')
+
+        branch3x3dbl = self.conv_bn_layer(
+            x, 64, 1, name=name + '.branch3x3dbl_1')
+        branch3x3dbl = self.conv_bn_layer(
+            branch3x3dbl, 96, 3, padding=1, name=name + '.branch3x3dbl_2')
+        branch3x3dbl = self.conv_bn_layer(
+            branch3x3dbl, 96, 3, padding=1, name=name + '.branch3x3dbl_3')
+
+        branch_pool = fluid.layers.pool2d(
+            x, pool_size=3, pool_stride=1, pool_padding=1, pool_type='avg')
+        branch_pool = self.conv_bn_layer(
+            branch_pool, pool_features, 1, name=name + '.branch_pool')
+
+        return fluid.layers.concat(
+            [branch1x1, branch5x5, branch3x3dbl, branch_pool], axis=1)
+
+    def inceptionB(self, x, name=None):
+        branch3x3 = self.conv_bn_layer(
+            x, 384, 3, stride=2, name=name + '.branch3x3')
+
+        branch3x3dbl = self.conv_bn_layer(
+            x, 64, 1, name=name + '.branch3x3dbl_1')
+        branch3x3dbl = self.conv_bn_layer(
+            branch3x3dbl, 96, 3, padding=1, name=name + '.branch3x3dbl_2')
+        branch3x3dbl = self.conv_bn_layer(
+            branch3x3dbl, 96, 3, stride=2, name=name + '.branch3x3dbl_3')
+
+        branch_pool = fluid.layers.pool2d(
+            x, pool_size=3, pool_stride=2, pool_type='max')
+
+        return fluid.layers.concat(
+            [branch3x3, branch3x3dbl, branch_pool], axis=1)
+
+    def inceptionC(self, x, c7, name=None):
+        branch1x1 = self.conv_bn_layer(x, 192, 1, name=name + '.branch1x1')
+
+        branch7x7 = self.conv_bn_layer(x, c7, 1, name=name + '.branch7x7_1')
+        branch7x7 = self.conv_bn_layer(
+            branch7x7, c7, (1, 7), padding=(0, 3), name=name + '.branch7x7_2')
+        branch7x7 = self.conv_bn_layer(
+            branch7x7, 192, (7, 1), padding=(3, 0), name=name + '.branch7x7_3')
+
+        branch7x7dbl = self.conv_bn_layer(
+            x, c7, 1, name=name + '.branch7x7dbl_1')
+        branch7x7dbl = self.conv_bn_layer(
+            branch7x7dbl,
+            c7, (7, 1),
+            padding=(3, 0),
+            name=name + '.branch7x7dbl_2')
+        branch7x7dbl = self.conv_bn_layer(
+            branch7x7dbl,
+            c7, (1, 7),
+            padding=(0, 3),
+            name=name + '.branch7x7dbl_3')
+        branch7x7dbl = self.conv_bn_layer(
+            branch7x7dbl,
+            c7, (7, 1),
+            padding=(3, 0),
+            name=name + '.branch7x7dbl_4')
+        branch7x7dbl = self.conv_bn_layer(
+            branch7x7dbl,
+            192, (1, 7),
+            padding=(0, 3),
+            name=name + '.branch7x7dbl_5')
+
+        branch_pool = fluid.layers.pool2d(
+            x, pool_size=3, pool_stride=1, pool_padding=1, pool_type='avg')
+        branch_pool = self.conv_bn_layer(
+            branch_pool, 192, 1, name=name + '.branch_pool')
+
+        return fluid.layers.concat(
+            [branch1x1, branch7x7, branch7x7dbl, branch_pool], axis=1)
+
+    def inceptionD(self, x, name=None):
+        branch3x3 = self.conv_bn_layer(x, 192, 1, name=name + '.branch3x3_1')
+        branch3x3 = self.conv_bn_layer(
+            branch3x3, 320, 3, stride=2, name=name + '.branch3x3_2')
+
+        branch7x7x3 = self.conv_bn_layer(
+            x, 192, 1, name=name + '.branch7x7x3_1')
+        branch7x7x3 = self.conv_bn_layer(
+            branch7x7x3,
+            192, (1, 7),
+            padding=(0, 3),
+            name=name + '.branch7x7x3_2')
+        branch7x7x3 = self.conv_bn_layer(
+            branch7x7x3,
+            192, (7, 1),
+            padding=(3, 0),
+            name=name + '.branch7x7x3_3')
+        branch7x7x3 = self.conv_bn_layer(
+            branch7x7x3, 192, 3, stride=2, name=name + '.branch7x7x3_4')
+
+        branch_pool = fluid.layers.pool2d(
+            x, pool_size=3, pool_stride=2, pool_type='max')
+        return fluid.layers.concat(
+            [branch3x3, branch7x7x3, branch_pool], axis=1)
+
+    def inceptionE(self, x, name=None):
+        branch1x1 = self.conv_bn_layer(x, 320, 1, name=name + '.branch1x1')
+
+        branch3x3 = self.conv_bn_layer(x, 384, 1, name=name + '.branch3x3_1')
+        branch3x3_2a = self.conv_bn_layer(
+            branch3x3,
+            384, (1, 3),
+            padding=(0, 1),
+            name=name + '.branch3x3_2a')
+        branch3x3_2b = self.conv_bn_layer(
+            branch3x3,
+            384, (3, 1),
+            padding=(1, 0),
+            name=name + '.branch3x3_2b')
+        branch3x3 = fluid.layers.concat([branch3x3_2a, branch3x3_2b], axis=1)
+
+        branch3x3dbl = self.conv_bn_layer(
+            x, 448, 1, name=name + '.branch3x3dbl_1')
+        branch3x3dbl = self.conv_bn_layer(
+            branch3x3dbl, 384, 3, padding=1, name=name + '.branch3x3dbl_2')
+        branch3x3dbl_3a = self.conv_bn_layer(
+            branch3x3dbl,
+            384, (1, 3),
+            padding=(0, 1),
+            name=name + '.branch3x3dbl_3a')
+        branch3x3dbl_3b = self.conv_bn_layer(
+            branch3x3dbl,
+            384, (3, 1),
+            padding=(1, 0),
+            name=name + '.branch3x3dbl_3b')
+        branch3x3dbl = fluid.layers.concat(
+            [branch3x3dbl_3a, branch3x3dbl_3b], axis=1)
+
+        branch_pool = fluid.layers.pool2d(
+            x, pool_size=3, pool_stride=1, pool_padding=1, pool_type='avg')
+        branch_pool = self.conv_bn_layer(
+            branch_pool, 192, 1, name=name + '.branch_pool')
+
+        return fluid.layers.concat(
+            [branch1x1, branch3x3, branch3x3dbl, branch_pool], axis=1)
+
+    def inceptionAux(self, x, num_classes, name=None):
+        x = fluid.layers.pool2d(x, pool_size=5, pool_stride=3, pool_type='avg')
+        x = self.conv_bn_layer(x, 128, 1, name=name + '.conv0')
+        x = self.conv_bn_layer(x, 768, 5, name=name + '.conv1')
+        x = fluid.layers.pool2d(x, global_pooling=True, pool_type='avg')
+        x = fluid.layers.flatten(x, axis=1)
+        x = fluid.layers.fc(x, size=num_classes)
+        return x
+
+    def fid_inceptionA(self, x, pool_features, name=None):
+        """ FID block in inception v3
+        """
+        branch1x1 = self.conv_bn_layer(x, 64, 1, name=name + '.branch1x1')
+
+        branch5x5 = self.conv_bn_layer(x, 48, 1, name=name + '.branch5x5_1')
+        branch5x5 = self.conv_bn_layer(
+            branch5x5, 64, 5, padding=2, name=name + '.branch5x5_2')
+
+        branch3x3dbl = self.conv_bn_layer(
+            x, 64, 1, name=name + '.branch3x3dbl_1')
+        branch3x3dbl = self.conv_bn_layer(
+            branch3x3dbl, 96, 3, padding=1, name=name + '.branch3x3dbl_2')
+        branch3x3dbl = self.conv_bn_layer(
+            branch3x3dbl, 96, 3, padding=1, name=name + '.branch3x3dbl_3')
+
+        branch_pool = fluid.layers.pool2d(
+            x,
+            pool_size=3,
+            pool_stride=1,
+            pool_padding=1,
+            exclusive=True,
+            pool_type='avg')
+        branch_pool = self.conv_bn_layer(
+            branch_pool, pool_features, 1, name=name + '.branch_pool')
+
+        return fluid.layers.concat(
+            [branch1x1, branch5x5, branch3x3dbl, branch_pool], axis=1)
+
+    def fid_inceptionC(self, x, c7, name=None):
+        """ FID block in inception v3
+        """
+        branch1x1 = self.conv_bn_layer(x, 192, 1, name=name + '.branch1x1')
+
+        branch7x7 = self.conv_bn_layer(x, c7, 1, name=name + '.branch7x7_1')
+        branch7x7 = self.conv_bn_layer(
+            branch7x7, c7, (1, 7), padding=(0, 3), name=name + '.branch7x7_2')
+        branch7x7 = self.conv_bn_layer(
+            branch7x7, 192, (7, 1), padding=(3, 0), name=name + '.branch7x7_3')
+
+        branch7x7dbl = self.conv_bn_layer(
+            x, c7, 1, name=name + '.branch7x7dbl_1')
+        branch7x7dbl = self.conv_bn_layer(
+            branch7x7dbl,
+            c7, (7, 1),
+            padding=(3, 0),
+            name=name + '.branch7x7dbl_2')
+        branch7x7dbl = self.conv_bn_layer(
+            branch7x7dbl,
+            c7, (1, 7),
+            padding=(0, 3),
+            name=name + '.branch7x7dbl_3')
+        branch7x7dbl = self.conv_bn_layer(
+            branch7x7dbl,
+            c7, (7, 1),
+            padding=(3, 0),
+            name=name + '.branch7x7dbl_4')
+        branch7x7dbl = self.conv_bn_layer(
+            branch7x7dbl,
+            192, (1, 7),
+            padding=(0, 3),
+            name=name + '.branch7x7dbl_5')
+
+        branch_pool = fluid.layers.pool2d(
+            x,
+            pool_size=3,
+            pool_stride=1,
+            pool_padding=1,
+            exclusive=True,
+            pool_type='avg')
+        branch_pool = self.conv_bn_layer(
+            branch_pool, 192, 1, name=name + '.branch_pool')
+
+        return fluid.layers.concat(
+            [branch1x1, branch7x7, branch7x7dbl, branch_pool], axis=1)
+
+    def fid_inceptionE_1(self, x, name=None):
+        """ FID block in inception v3
+        """
+        branch1x1 = self.conv_bn_layer(x, 320, 1, name=name + '.branch1x1')
+
+        branch3x3 = self.conv_bn_layer(x, 384, 1, name=name + '.branch3x3_1')
+        branch3x3_2a = self.conv_bn_layer(
+            branch3x3,
+            384, (1, 3),
+            padding=(0, 1),
+            name=name + '.branch3x3_2a')
+        branch3x3_2b = self.conv_bn_layer(
+            branch3x3,
+            384, (3, 1),
+            padding=(1, 0),
+            name=name + '.branch3x3_2b')
+        branch3x3 = fluid.layers.concat([branch3x3_2a, branch3x3_2b], axis=1)
+
+        branch3x3dbl = self.conv_bn_layer(
+            x, 448, 1, name=name + '.branch3x3dbl_1')
+        branch3x3dbl = self.conv_bn_layer(
+            branch3x3dbl, 384, 3, padding=1, name=name + '.branch3x3dbl_2')
+        branch3x3dbl_3a = self.conv_bn_layer(
+            branch3x3dbl,
+            384, (1, 3),
+            padding=(0, 1),
+            name=name + '.branch3x3dbl_3a')
+        branch3x3dbl_3b = self.conv_bn_layer(
+            branch3x3dbl,
+            384, (3, 1),
+            padding=(1, 0),
+            name=name + '.branch3x3dbl_3b')
+        branch3x3dbl = fluid.layers.concat(
+            [branch3x3dbl_3a, branch3x3dbl_3b], axis=1)
+
+        branch_pool = fluid.layers.pool2d(
+            x,
+            pool_size=3,
+            pool_stride=1,
+            pool_padding=1,
+            exclusive=True,
+            pool_type='avg')
+        branch_pool = self.conv_bn_layer(
+            branch_pool, 192, 1, name=name + '.branch_pool')
+
+        return fluid.layers.concat(
+            [branch1x1, branch3x3, branch3x3dbl, branch_pool], axis=1)
+
+    def fid_inceptionE_2(self, x, name=None):
+        """ FID block in inception v3
+        """
+        branch1x1 = self.conv_bn_layer(x, 320, 1, name=name + '.branch1x1')
+
+        branch3x3 = self.conv_bn_layer(x, 384, 1, name=name + '.branch3x3_1')
+        branch3x3_2a = self.conv_bn_layer(
+            branch3x3,
+            384, (1, 3),
+            padding=(0, 1),
+            name=name + '.branch3x3_2a')
+        branch3x3_2b = self.conv_bn_layer(
+            branch3x3,
+            384, (3, 1),
+            padding=(1, 0),
+            name=name + '.branch3x3_2b')
+        branch3x3 = fluid.layers.concat([branch3x3_2a, branch3x3_2b], axis=1)
+
+        branch3x3dbl = self.conv_bn_layer(
+            x, 448, 1, name=name + '.branch3x3dbl_1')
+        branch3x3dbl = self.conv_bn_layer(
+            branch3x3dbl, 384, 3, padding=1, name=name + '.branch3x3dbl_2')
+        branch3x3dbl_3a = self.conv_bn_layer(
+            branch3x3dbl,
+            384, (1, 3),
+            padding=(0, 1),
+            name=name + '.branch3x3dbl_3a')
+        branch3x3dbl_3b = self.conv_bn_layer(
+            branch3x3dbl,
+            384, (3, 1),
+            padding=(1, 0),
+            name=name + '.branch3x3dbl_3b')
+        branch3x3dbl = fluid.layers.concat(
+            [branch3x3dbl_3a, branch3x3dbl_3b], axis=1)
+
+        ### same with paper
+        branch_pool = fluid.layers.pool2d(
+            x, pool_size=3, pool_stride=1, pool_padding=1, pool_type='max')
+        branch_pool = self.conv_bn_layer(
+            branch_pool, 192, 1, name=name + '.branch_pool')
+
+        return fluid.layers.concat(
+            [branch1x1, branch3x3, branch3x3dbl, branch_pool], axis=1)
+
+    def conv_bn_layer(self,
+                      data,
+                      num_filters,
+                      filter_size,
+                      stride=1,
+                      padding=0,
+                      groups=1,
+                      act='relu',
+                      name=None):
+        conv = fluid.layers.conv2d(
+            input=data,
+            num_filters=num_filters,
+            filter_size=filter_size,
+            stride=stride,
+            padding=padding,
+            groups=groups,
+            act=None,
+            param_attr=ParamAttr(name=name + ".conv.weight"),
+            bias_attr=False,
+            name=name)
+        return fluid.layers.batch_norm(
+            input=conv,
+            act=act,
+            epsilon=0.001,
+            name=name + '.bn',
+            param_attr=ParamAttr(name=name + ".bn.weight"),
+            bias_attr=ParamAttr(name=name + ".bn.bias"),
+            moving_mean_name=name + '.bn.running_mean',
+            moving_variance_name=name + '.bn.running_var')

demo/gan_compression/metric/params_inceptionV3

+../../../Paddle-GAN-compression/metric/params_inceptionV3/
\ No newline at end of file

demo/gan_compression/models/__init__.py

+import importlib
+from .modules import *
+from .base_model import BaseModel
+
+
+def find_model_using_name(model_name):
+    model_filename = "models." + model_name + "_model"
+    modellib = importlib.import_module(model_filename)
+    target_model_name = model_name.replace('_', '')
+    model = None
+    for name, cls in modellib.__dict__.items():
+        if name.lower() == target_model_name.lower() and issubclass(cls,
+                                                                    BaseModel):
+            model = cls
+    assert model is not None, "model {} is not right, please check it!".format(
+        model_name)
+
+    return model
+
+
+def get_special_cfg(model):
+    model_cls = find_model_using_name(model)
+    return model_cls.add_special_cfgs
+
+
+def create_model(cfg):
+    model_cls = find_model_using_name(cfg.model)
+    return model_cls(cfg)

demo/gan_compression/models/base_model.py

+import os
+import paddle.fluid as fluid
+
+
+class BaseModel(fluid.dygraph.Layer):
+    @staticmethod
+    def add_special_cfgs(parser):
+        pass
+
+    def set_input(self, inputs):
+        pass
+
+    def setup(self):
+        self.load_network()
+
+    def load_network(self):
+        for name in self.model_names:
+            net = getattr(self, 'net' + name, None)
+            path = getattr(self.args, 'restore_%s_path' % name, None)
+            if path is not None:
+                util.load_network(net, path)
+
+    def save_network(self, epoch):
+        for name in self.model_names:
+            if isinstance(name, str):
+                save_filename = '%s_net_%s' % (epoch, name)
+                save_path = os.path.join(self.args.save_dir, save_filename)
+                net = getattr(self, 'net' + name)
+                fluid.save_dygraph(net.state_dict(), save_path)
+
+    def forward(self):
+        pass
+
+    def optimize_parameter(self):
+        pass
+
+    def get_current_loss(self):
+        loss_dict = {}
+        for name in self.loss_names:
+            if not hasattr(self, 'loss_' + name):
+                continue
+            key = name
+            loss_dict[key] = float(getattr(self, 'loss_' + name))
+        return loss_dict
+
+    def get_current_lr(self):
+        raise NotImplementedError
+
+    def set_stop_gradient(self, nets, stop_grad=False):
+        if not isinstance(nets, list):
+            nets = [nets]
+        for net in nets:
+            if net is not None:
+                for param in net.parameters():
+                    param.stop_gradient = stop_grad
+
+    def evaluate_model(self):
+        pass
+
+    def profile(self):
+        pass

demo/gan_compression/models/cycle_gan_model.py

+import itertools
+import os
+import numpy as np
+import paddle.fluid as fluid
+from dataset.data_loader import create_eval_data
+from utils.image_pool import ImagePool
+from models import network, loss
+from models.base_model import BaseModel
+from metric.inception import InceptionV3
+from utils import util, optimization
+from metric import get_fid
+
+
+class CycleGAN(BaseModel):
+    @staticmethod
+    def add_special_cfgs(parser):
+        parser.add_argument(
+            '--mobile_lr',
+            type=float,
+            default=2e-4,
+            help="initial learning rate to train cyclegan")
+        parser.add_argument(
+            '--mobile_epoch',
+            type=int,
+            default=200,
+            help="The number of epoch to train mobile net")
+        parser.add_argument(
+            '--mobile_nepochs',
+            type=int,
+            default=100,
+            help="number of epochs with the initial learning rate")
+        parser.add_argument(
+            '--mobile_nepochs_decay',
+            type=int,
+            default=100,
+            help="number of epochs to linearly decay learning rate to zero")
+        parser.add_argument(
+            '--mobile_scheduler',
+            type=str,
+            default='linear',
+            help="learning rate scheduler")
+        parser.add_argument(
+            '--real_stat_A_path',
+            type=str,
+            default='real_stat/horse2zebra_A.npz',
+            help="")
+        parser.add_argument(
+            '--real_stat_B_path',
+            type=str,
+            default='real_stat/horse2zebra_B.npz',
+            help="")
+        parser.add_argument(
+            '--recon_loss_mode',
+            type=str,
+            default='l1',
+            choices=['l1', 'l2'],
+            help="")
+        parser.add_argument(
+            '--lambda_A',
+            type=float,
+            default=10.0,
+            help="weight to scale cycle loss (A->B->A)")
+        parser.add_argument(
+            '--lambda_B',
+            type=float,
+            default=10.0,
+            help="weight to scale cycle loss (B->A->B)")
+        parser.add_argument(
+            '--lambda_identity',
+            type=float,
+            default=0.5,
+            help="Weight to scale identity mapping loss")
+        parser.add_argument(
+            '--pool_size', type=int, default=50, help="pool size in cyclegan")
+        return parser
+
+    def __init__(self, cfgs):
+        super(CycleGAN, self).__init__()
+        assert cfgs.direction == 'AtoB'
+        self.cfgs = cfgs
+        self.loss_names = [
+            'D_A', 'G_A', 'G_cycle_A', 'G_idt_A', 'D_B', 'G_B', 'G_cycle_B',
+            'G_idt_B'
+        ]
+        self.model_names = ['G_A', 'G_B', 'D_A', 'D_B']
+
+        self.netG_A = network.define_G(cfgs.input_nc, cfgs.output_nc, cfgs.ngf,
+                                       cfgs.netG, cfgs.norm_type,
+                                       cfgs.dropout_rate)
+        self.netG_B = network.define_G(cfgs.output_nc, cfgs.input_nc, cfgs.ngf,
+                                       cfgs.netG, cfgs.norm_type,
+                                       cfgs.dropout_rate)
+        self.netD_A = network.define_D(cfgs.output_nc, cfgs.ndf, cfgs.netD,
+                                       cfgs.norm_type, cfgs.n_layer_D)
+        self.netD_B = network.define_D(cfgs.input_nc, cfgs.ndf, cfgs.netD,
+                                       cfgs.norm_type, cfgs.n_layer_D)
+
+        if cfgs.lambda_identity > 0.0:
+            assert (cfgs.input_nc == cfgs.output_nc)
+        self.fake_A_pool = ImagePool(cfgs.pool_size)
+        self.fake_B_pool = ImagePool(cfgs.pool_size)
+
+        self.optimizer_G = optimization.Optimizer(
+            cfgs.mobile_lr,
+            cfgs.mobile_scheduler,
+            cfgs.step_per_epoch,
+            cfgs.mobile_nepochs,
+            cfgs.mobile_nepochs_decay,
+            cfgs,
+            parameter_list=(
+                self.netG_A.parameters() + self.netG_B.parameters()))
+        self.optimizer_D_A = optimization.Optimizer(
+            cfgs.mobile_lr,
+            cfgs.mobile_scheduler,
+            cfgs.step_per_epoch,
+            cfgs.mobile_nepochs,
+            cfgs.mobile_nepochs_decay,
+            cfgs,
+            parameter_list=self.netD_A.parameters())
+        self.optimizer_D_B = optimization.Optimizer(
+            cfgs.mobile_lr,
+            cfgs.mobile_scheduler,
+            cfgs.step_per_epoch,
+            cfgs.mobile_nepochs,
+            cfgs.mobile_nepochs_decay,
+            cfgs,
+            parameter_list=self.netD_B.parameters())
+
+        self.eval_dataloader_AtoB, self.name_AtoB = create_eval_data(
+            cfgs, direction='AtoB')
+        self.eval_dataloader_BtoA, self.name_BtoA = create_eval_data(
+            cfgs, direction='BtoA')
+
+        block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[2048]
+        self.inception_model = InceptionV3([block_idx])
+
+        self.best_fid_A, self.best_fid_B = 1e9, 1e9
+        self.fids_A, self.fids_B = [], []
+        self.is_best = False
+        self.npz_A = np.load(cfgs.real_stat_A_path)
+        self.npz_B = np.load(cfgs.real_stat_B_path)
+
+    def set_input(self, inputs):
+        self.real_A = inputs[0] if self.cfgs.direction == 'AtoB' else inputs[1]
+        self.real_B = inputs[1] if self.cfgs.direction == 'AtoB' else inputs[0]
+
+    def set_single_input(self, inputs):
+        self.real_A = inputs[0]
+
+    def setup(self):
+        self.load_network()
+
+    def load_network(self):
+        for name in self.model_names:
+            net = getattr(self, 'net' + name, None)
+            path = getattr(self.cfgs, 'restore_%s_path' % name, None)
+            if path is not None:
+                util.load_network(net, path)
+
+    def save_network(self, epoch):
+        for name in self.model_names:
+            if isinstance(name, str):
+                save_filename = '%s_net%s' % (epoch, name)
+                save_path = os.path.join(self.cfgs.save_dir, 'mobile',
+                                         save_filename)
+                net = getattr(self, 'net' + name)
+                fluid.save_dygraph(net.state_dict(), save_path)
+
+    def forward(self):
+        self.fake_B = self.netG_A(self.real_A)  ## G_A(A)
+        self.rec_A = self.netG_B(self.fake_B)  ## G_B(G_A(A))
+        self.fake_A = self.netG_B(self.real_B)  ## G_B(B)
+        self.rec_B = self.netG_A(self.fake_A)  ## G_A(G_B(B))
+
+    def backward_D_basic(self, netD, real, fake):
+        ### real
+        pred_real = netD(real)
+        loss_D_real = loss.gan_loss(self.cfgs.gan_loss_mode, pred_real, True)
+        ### fake
+        pred_fake = netD(fake.detach())
+        loss_D_fake = loss.gan_loss(self.cfgs.gan_loss_mode, pred_fake, False)
+
+        loss_D = (loss_D_real + loss_D_fake) * 0.5
+        loss_D.backward()
+        return loss_D
+
+    def backward_D_A(self):
+        fake_B = self.fake_B_pool.query(self.fake_B)
+        self.loss_D_A = self.backward_D_basic(self.netD_A, self.real_B, fake_B)
+
+    def backward_D_B(self):
+        fake_A = self.fake_A_pool.query(self.fake_A)
+        self.loss_D_B = self.backward_D_basic(self.netD_B, self.real_A, fake_A)
+
+    def backward_G(self):
+        lambda_idt = self.cfgs.lambda_identity
+        lambda_A = self.cfgs.lambda_A
+        lambda_B = self.cfgs.lambda_B
+
+        if lambda_idt > 0:
+            ### identity loss G_A: ||G_A(B) - B||
+            self.idt_A = self.netG_A(self.real_B)
+            self.loss_G_idt_A = loss.recon_loss(
+                'l1', self.idt_A, self.real_B) * lambda_B * lambda_idt
+            ### identity loss G_B: ||G_B(A) - A||
+            self.idt_B = self.netG_B(self.real_A)
+            self.loss_G_idt_B = loss.recon_loss(
+                'l1', self.idt_B, self.real_A) * lambda_A * lambda_idt
+
+        else:
+            self.loss_G_idt_A = 0
+            self.loss_G_idt_B = 0
+
+        ### GAN loss D_A(G_A(A))
+        self.loss_G_A = loss.gan_loss(self.cfgs.gan_loss_mode,
+                                      self.netD_A(self.fake_B), True)
+        ### GAN loss D_B(G_B(B))
+        self.loss_G_B = loss.gan_loss(self.cfgs.gan_loss_mode,
+                                      self.netD_B(self.fake_A), True)
+        ### forward cycle loss ||G_B(G_A(A)) - A||
+        self.loss_G_cycle_A = loss.recon_loss('l1', self.rec_A,
+                                              self.real_A) * lambda_A
+        ### backward cycle loss ||G_A(G_B(B)) - B||
+        self.loss_G_cycle_B = loss.recon_loss('l1', self.rec_B,
+                                              self.real_B) * lambda_B
+        ### combine loss and calculate gradients
+        self.loss_G = self.loss_G_A + self.loss_G_B + self.loss_G_cycle_A + self.loss_G_cycle_B + self.loss_G_idt_A + self.loss_G_idt_B
+
+        self.loss_G.backward()
+
+    def optimize_parameter(self):
+        self.forward()
+        self.set_stop_gradient([self.netD_A, self.netD_B], True)
+        self.backward_G()  ## calculate gradients for G_A and G_B
+        self.optimizer_G.optimizer.minimize(self.loss_G)
+        self.optimizer_G.optimizer.clear_gradients()
+
+        self.set_stop_gradient([self.netD_A], False)
+        self.backward_D_A()  ### calculate gradients for D_A
+        self.optimizer_D_A.optimizer.minimize(self.loss_D_A)
+        self.optimizer_D_A.optimizer.clear_gradients()
+
+        self.set_stop_gradient([self.netD_B], False)
+        self.backward_D_B()  ### calculate gradients for D_B
+        self.optimizer_D_B.optimizer.minimize(self.loss_D_B)
+        self.optimizer_D_B.optimizer.clear_gradients()
+
+    @fluid.dygraph.no_grad
+    def test_single_side(self, direction):
+        generator = getattr(self, 'netG_%s' % direction[0])
+        self.fake_B = generator(self.real_A)
+
+    def get_current_lr(self):
+        lr_dict = {}
+        lr_dict['optim_G'] = self.optimizer_G.optimizer.current_step_lr()
+        lr_dict['optim_D_A'] = self.optimizer_D_A.optimizer.current_step_lr()
+        lr_dict['optim_D_B'] = self.optimizer_D_B.optimizer.current_step_lr()
+        return lr_dict
+
+    def evaluate_model(self, step):
+        ret = {}
+        self.is_best = False
+        save_dir = os.path.join(self.cfgs.save_dir, 'mobile', 'eval',
+                                str(step))
+        if not os.path.exists(save_dir):
+            os.makedirs(save_dir)
+        self.netG_A.eval()
+        self.netG_B.eval()
+        for direction in ['AtoB', 'BtoA']:
+            eval_dataloader = getattr(self, 'eval_dataloader_' + direction)
+            id2name = getattr(self, 'name_' + direction)
+            fakes = []
+            cnt = 0
+            for i, data_i in enumerate(eval_dataloader):
+                self.set_single_input(data_i)
+                self.test_single_side(direction)
+                fakes.append(self.fake_B.detach().numpy())
+                for j in range(len(self.fake_B)):
+                    if cnt < 10:
+                        name = 'fake_' + direction + str(id2name[i +
+                                                                 j]) + '.png'
+                        save_path = os.path.join(save_dir, name)
+                        fake_im = util.tensor2img(self.fake_B[j])
+                        util.save_image(fake_im, save_path)
+                    cnt += 1
+
+            suffix = direction[-1]
+            fluid.disable_imperative()
+            fid = get_fid(fakes, self.inception_model,
+                          getattr(self, 'npz_%s' % direction[-1]),
+                          self.cfgs.inception_model)
+            fluid.enable_imperative()
+            if fid < getattr(self, 'best_fid_%s' % suffix):
+                self.is_best = True
+                setattr(self, 'best_fid_%s' % suffix, fid)
+            print("direction: %s, fid score is: %f, best fid score is %f" %
+                  (direction, fid, getattr(self, 'best_fid_%s' % suffix)))
+            fids = getattr(self, 'fids_%s' % suffix)
+            fids.append(fid)
+            if len(fids) > 3:
+                fids.pop(0)
+            ret['metric/fid_%s' % suffix] = fid
+            ret['metric/fid_%s-mean' %
+                suffix] = sum(getattr(self, 'fids_%s' % suffix)) / len(
+                    getattr(self, 'fids_%s' % suffix))
+            ret['metric/fid_%s-best' % suffix] = getattr(self, 'best_fid_%s' %
+                                                         suffix)
+
+        self.netG_A.train()
+        self.netG_B.train()
+        return ret

demo/gan_compression/models/discrimitor.py

+import functools
+import paddle.fluid as fluid
+from paddle.fluid.dygraph.nn import InstanceNorm, Conv2D, Conv2DTranspose, BatchNorm
+from paddle.nn.layer import Leaky_ReLU, ReLU, Pad2D
+
+
+class NLayerDiscriminator(fluid.dygraph.Layer):
+    def __init__(self, input_channel, ndf, n_layers=3,
+                 norm_layer=InstanceNorm):
+        super(NLayerDiscriminator, self).__init__()
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == InstanceNorm
+        else:
+            use_bias = norm_layer == InstanceNorm
+
+        kw = 4
+        padw = 1
+        self.model = fluid.dygraph.LayerList([
+            Conv2D(
+                input_channel, ndf, filter_size=kw, stride=2, padding=padw),
+            Leaky_ReLU(0.2)
+        ])
+        nf_mult = 1
+        nf_mult_prev = 1
+        for n in range(1, n_layers):
+            nf_mult_prev = nf_mult
+            nf_mult = min(2**n, 8)
+            self.model.extend([
+                Conv2D(
+                    ndf * nf_mult_prev,
+                    ndf * nf_mult,
+                    filter_size=kw,
+                    stride=2,
+                    padding=padw,
+                    bias_attr=use_bias),
+                #norm_layer(ndf * nf_mult),
+                InstanceNorm(
+                    ndf * nf_mult,
+                    param_attr=fluid.ParamAttr(
+                        initializer=fluid.initializer.Constant(1.0),
+                        learning_rate=0.0,
+                        trainable=False),
+                    bias_attr=fluid.ParamAttr(
+                        initializer=fluid.initializer.Constant(0.0),
+                        learning_rate=0.0,
+                        trainable=False)),
+                Leaky_ReLU(0.2)
+            ])
+
+        nf_mult_prev = nf_mult
+        nf_mult = min(2**n_layers, 8)
+        self.model.extend([
+            Conv2D(
+                ndf * nf_mult_prev,
+                ndf * nf_mult,
+                filter_size=kw,
+                stride=1,
+                padding=padw,
+                bias_attr=use_bias),
+            #norm_layer(ndf * nf_mult),
+            InstanceNorm(
+                ndf * nf_mult,
+                param_attr=fluid.ParamAttr(
+                    initializer=fluid.initializer.Constant(1.0),
+                    learning_rate=0.0,
+                    trainable=False),
+                bias_attr=fluid.ParamAttr(
+                    initializer=fluid.initializer.Constant(0.0),
+                    learning_rate=0.0,
+                    trainable=False)),
+            Leaky_ReLU(0.2)
+        ])
+
+        self.model.extend([
+            Conv2D(
+                ndf * nf_mult, 1, filter_size=kw, stride=1, padding=padw)
+        ])
+
+    def forward(self, inputs):
+        #import numpy as np
+        #print("================ DISCRIMINATOR ====================")
+        y = inputs
+        for sublayer in self.model:
+            y = sublayer(y)
+        #    print(sublayer, np.sum(np.abs(y.numpy())))
+        #print("===================================================")
+        return y

demo/gan_compression/models/generator/mobile_generator.py

+import functools
+import paddle.fluid as fluid
+from paddle.fluid.dygraph.nn import InstanceNorm, Conv2D, Conv2DTranspose
+from paddle.nn.layer import Leaky_ReLU, ReLU, Pad2D
+from ..modules import MobileResnetBlock
+
+use_cudnn = False
+
+
+class MobileResnetGenerator(fluid.dygraph.Layer):
+    def __init__(self,
+                 input_channel,
+                 output_nc,
+                 ngf,
+                 norm_layer=InstanceNorm,
+                 dropout_rate=0,
+                 n_blocks=9,
+                 padding_type='reflect'):
+        super(MobileResnetGenerator, self).__init__()
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == InstanceNorm
+        else:
+            use_bias = norm_layer == InstanceNorm
+
+        self.model = fluid.dygraph.LayerList([
+            Pad2D(
+                paddings=[3, 3, 3, 3], mode="reflect"), Conv2D(
+                    input_channel,
+                    int(ngf),
+                    filter_size=7,
+                    padding=0,
+                    use_cudnn=use_cudnn,
+                    bias_attr=use_bias), norm_layer(ngf), ReLU()
+        ])
+
+        n_downsampling = 2
+        for i in range(n_downsampling):
+            mult = 2**i
+            self.model.extend([
+                Conv2D(
+                    ngf * mult,
+                    ngf * mult * 2,
+                    filter_size=3,
+                    stride=2,
+                    padding=1,
+                    use_cudnn=use_cudnn,
+                    bias_attr=use_bias), norm_layer(ngf * mult * 2), ReLU()
+            ])
+
+        mult = 2**n_downsampling
+
+        n_blocks1 = n_blocks // 3
+        n_blocks2 = n_blocks1
+        n_blocks3 = n_blocks - n_blocks1 - n_blocks2
+
+        for i in range(n_blocks1):
+            self.model.extend([
+                MobileResnetBlock(
+                    ngf * mult,
+                    ngf * mult,
+                    padding_type=padding_type,
+                    norm_layer=norm_layer,
+                    dropout_rate=dropout_rate,
+                    use_bias=use_bias)
+            ])
+
+        for i in range(n_blocks2):
+            self.model.extend([
+                MobileResnetBlock(
+                    ngf * mult,
+                    ngf * mult,
+                    padding_type=padding_type,
+                    norm_layer=norm_layer,
+                    dropout_rate=dropout_rate,
+                    use_bias=use_bias)
+            ])
+
+        for i in range(n_blocks3):
+            self.model.extend([
+                MobileResnetBlock(
+                    ngf * mult,
+                    ngf * mult,
+                    padding_type=padding_type,
+                    norm_layer=norm_layer,
+                    dropout_rate=dropout_rate,
+                    use_bias=use_bias)
+            ])
+
+        for i in range(n_downsampling):
+            mult = 2**(n_downsampling - i)
+            output_size = (i + 1) * 128
+            self.model.extend([
+                Conv2DTranspose(
+                    ngf * mult,
+                    int(ngf * mult / 2),
+                    filter_size=3,
+                    output_size=output_size,
+                    stride=2,
+                    padding=1,
+                    use_cudnn=use_cudnn,
+                    bias_attr=use_bias), norm_layer(int(ngf * mult / 2)),
+                ReLU()
+            ])
+
+        self.model.extend([Pad2D(paddings=[3, 3, 3, 3], mode="reflect")])
+        self.model.extend([Conv2D(ngf, output_nc, filter_size=7, padding=0)])
+
+    def forward(self, inputs):
+        y = inputs
+        for sublayer in self.model:
+            y = sublayer(y)
+        y = fluid.layers.tanh(y)
+        return y

demo/gan_compression/models/generator/resnet_generator.py

+import functools
+import paddle.fluid as fluid
+from paddle.fluid.dygraph.nn import InstanceNorm, Conv2D, Conv2DTranspose
+from paddle.nn.layer import Leaky_ReLU, ReLU, Pad2D
+from ..modules import ResnetBlock
+
+
+class ResnetGenerator(fluid.dygraph.Layer):
+    def __init__(self,
+                 input_channel,
+                 output_nc,
+                 ngf,
+                 norm_layer=InstanceNorm,
+                 dropout_rate=0,
+                 n_blocks=6,
+                 padding_type='reflect'):
+        super(ResnetGenerator, self).__init__()
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == InstanceNorm
+        else:
+            use_bias = norm_layer == InstanceNorm
+
+        self.model = fluid.dygraph.LayerList([
+            Pad2D(
+                paddings=[3, 3, 3, 3], mode="reflect"), Conv2D(
+                    input_channel,
+                    int(ngf),
+                    filter_size=7,
+                    padding=0,
+                    bias_attr=use_bias), norm_layer(ngf), ReLU()
+        ])
+
+        n_downsampling = 2
+        for i in range(n_downsampling):
+            mult = 2**i
+            self.model.extend([
+                Conv2D(
+                    ngf * mult,
+                    ngf * mult * 2,
+                    filter_size=3,
+                    stride=2,
+                    padding=1,
+                    bias_attr=use_bias), norm_layer(ngf * mult * 2), ReLU()
+            ])
+
+        mult = 2**n_downsampling
+        for i in range(n_blocks):
+            self.model.extend([
+                ResnetBlock(
+                    ngf * mult,
+                    padding_type=padding_type,
+                    norm_layer=norm_layer,
+                    dropout_rate=dropout_rate,
+                    use_bias=use_bias)
+            ])
+
+        for i in range(n_downsampling):
+            mult = 2**(n_downsampling - i)
+            self.model.extend([
+                Conv2DTranspose(
+                    ngf * mult,
+                    int(ngf * mult / 2),
+                    filter_size=3,
+                    stride=2,
+                    padding=1,
+                    bias_attr=use_bias), Pad2D(
+                        paddings=[0, 1, 0, 1], mode='constant', pad_value=0.0),
+                norm_layer(int(ngf * mult / 2)), ReLU()
+            ])
+
+        self.model.extend([Pad2D(paddings=[3, 3, 3, 3], mode="reflect")])
+        self.model.extend([Conv2D(ngf, output_nc, filter_size=7, padding=0)])
+
+    def forward(self, inputs):
+        y = fluid.layers.clamp(inputs, min=-1.0, max=1.0)
+        for sublayer in self.model:
+            y = sublayer(y)
+        y = fluid.layers.tanh(y)
+        return y

demo/gan_compression/models/generator/sub_mobile_generator.py

+import functools
+import paddle.fluid as fluid
+import paddle.tensor as tensor
+from paddle.fluid.dygraph.nn import InstanceNorm, Conv2D, Conv2DTranspose
+from paddle.nn.layer import Leaky_ReLU, ReLU, Pad2D
+from .modules import SeparableConv2D, MobileResnetBlock
+
+use_cudnn = False
+
+
+class SubMobileResnetGenerator(fluid.dygraph.Layer):
+    def __init__(self,
+                 input_channel,
+                 output_nc,
+                 config,
+                 norm_layer=InstanceNorm,
+                 dropout_rate=0,
+                 n_blocks=9,
+                 padding_type='reflect'):
+        super(SubMobileResnetGenerator, self).__init__()
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == InstanceNorm
+        else:
+            use_bias = norm_layer == InstanceNorm
+
+        self.model = fluid.dygraph.LayerList([
+            Pad2D(
+                paddings=[3, 3, 3, 3], mode="reflect"), Conv2D(
+                    input_channel,
+                    config['channels'][0],
+                    filter_size=7,
+                    padding=0,
+                    use_cudnn=use_cudnn,
+                    bias_attr=use_bias), norm_layer(config['channels'][0]),
+            ReLU()
+        ])
+
+        n_downsampling = 2
+        for i in range(n_downsampling):
+            mult = 2**i
+            in_c = config['channels'][i]
+            out_c = config['channels'][i + 1]
+            self.model.extend([
+                Conv2D(
+                    in_c * mult,
+                    out_c * mult * 2,
+                    filter_size=3,
+                    stride=2,
+                    padding=1,
+                    use_cudnn=use_cudnn,
+                    bias_attr=use_bias), norm_layer(out_c * mult * 2), ReLU()
+            ])
+
+        mult = 2**n_downsampling
+
+        in_c = config['channels'][2]
+        for i in range(n_blocks):
+            if len(config['channels']) == 6:
+                offset = 0
+            else:
+                offset = i // 3
+            out_c = config['channels'][offset + 3]
+            self.model.extend([
+                MobileResnetBlock(
+                    in_c * mult,
+                    out_c * mult,
+                    padding_type=padding_type,
+                    norm_layer=norm_layer,
+                    dropout_rate=dropout_rate,
+                    use_bias=use_bias)
+            ])
+
+        if len(config['channels']) == 6:
+            offset = 4
+        else:
+            offset = 6
+        for i in range(n_downsampling):
+            out_c = config['channels'][offset + i]
+            mult = 2**(n_downsampling - i)
+            output_size = (i + 1) * 128
+            self.model.extend([
+                Conv2DTranspose(
+                    in_c * mult,
+                    int(out_c * mult / 2),
+                    filter_size=3,
+                    output_size=output_size,
+                    stride=2,
+                    padding=1,
+                    use_cudnn=use_cudnn,
+                    bias_attr=use_bias), norm_layer(int(out_c * mult / 2)),
+                ReLU()
+            ])
+            in_c = out_c
+
+        self.model.extend([Pad2D(paddings=[3, 3, 3, 3], mode="reflect")])
+        self.model.extend([Conv2D(in_c, output_nc, filter_size=7, padding=0)])
+
+    def forward(self, inputs):
+        y = tensor.clamp(input, min=-1, max=1)
+        for sublayer in self.model:
+            y = sublayer(y)
+        y = fluid.layers.tanh(y)
+        return y

demo/gan_compression/models/generator/super_generator.py

+import functools
+import paddle.fluid as fluid
+import paddle.tensor as tensor
+from paddle.fluid.dygraph.nn import BatchNorm, InstanceNorm, Dropout
+from paddle.nn.layer import Leaky_ReLU, ReLU, Pad2D
+from ..super_modules import SuperConv2D, SuperConv2DTranspose, SuperSeparableConv2D, SuperInstanceNorm
+
+
+class SuperMobileResnetBlock(fluid.dygraph.Layer):
+    def __init__(self, dim, padding_type, norm_layer, dropout_rate, use_bias):
+        super(SuperMobileResnetBlock, self).__init__()
+        self.conv_block = fluid.dygraph.LayerList([])
+        p = 0
+        if padding_type == 'reflect':
+            self.conv_block.extend(
+                [Pad2D(
+                    paddings=[1, 1, 1, 1], mode="reflect")])
+        elif padding_type == 'replicate':
+            self.conv_block.extend([Pad2D(paddings=[1, 1, 1, 1], mode="edge")])
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' %
+                                      self.padding_type)
+
+        self.conv_block.extend([
+            SuperSeparableConv2D(
+                num_channels=dim,
+                num_filters=dim,
+                filter_size=3,
+                stride=1,
+                padding=p), norm_layer(dim), ReLU()
+        ])
+        self.conv_block.extend([Dropout(dropout_rate)])
+
+        p = 0
+        if padding_type == 'reflect':
+            self.conv_block.extend(
+                [Pad2D(
+                    paddings=[1, 1, 1, 1], mode="reflect")])
+        elif padding_type == 'replicate':
+            self.conv_block.extend([Pad2D(paddings=[1, 1, 1, 1], mode="edge")])
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' %
+                                      self.padding_type)
+
+        self.conv_block.extend([
+            SuperSeparableConv2D(
+                num_channels=dim,
+                num_filters=dim,
+                filter_size=3,
+                stride=1,
+                padding=p), norm_layer(dim)
+        ])
+
+    def forward(self, input, config):
+        x = input
+        cnt = 0
+        for sublayer in self.conv_block:
+            if isinstance(sublayer, SuperSeparableConv2D):
+                if cnt == 1:
+                    config['channel'] = input.shape[1]
+                x = sublayer(x, config)
+                cnt += 1
+            else:
+                x = sublayer(x)
+        out = input + x
+        return out
+
+
+class SuperMobileResnetGenerator(fluid.dygraph.Layer):
+    def __init__(self,
+                 input_channel,
+                 output_nc,
+                 ngf,
+                 norm_layer=BatchNorm,
+                 dropout_rate=0,
+                 n_blocks=6,
+                 padding_type='reflect'):
+        assert n_blocks >= 0
+        super(SuperMobileResnetGenerator, self).__init__()
+        if norm_layer.func == InstanceNorm or norm_layer == InstanceNorm:
+            norm_layer = SuperInstanceNorm
+        else:
+            raise NotImplementedError
+
+        use_bias = norm_layer == InstanceNorm
+
+        self.model = fluid.dygraph.LayerList([])
+        self.model.extend([
+            Pad2D(
+                paddings=[3, 3, 3, 3], mode="reflect"), SuperConv2D(
+                    input_channel,
+                    ngf,
+                    filter_size=7,
+                    padding=0,
+                    bias_attr=use_bias), norm_layer(ngf), ReLU()
+        ])
+
+        n_downsampling = 2
+        for i in range(n_downsampling):
+            mult = 2**i
+            self.model.extend([
+                SuperConv2D(
+                    ngf * mult,
+                    ngf * mult * 2,
+                    filter_size=3,
+                    stride=2,
+                    padding=1,
+                    bias_attr=use_bias), norm_layer(int(ngf * mult * 2)),
+                ReLU()
+            ])
+
+        mult = 2**n_downsampling
+        n_blocks1 = n_blocks // 3
+        n_blocks2 = n_blocks1
+        n_blocks3 = n_blocks - n_blocks1 - n_blocks2
+
+        for i in range(n_blocks1):
+            self.model.extend([
+                SuperMobileResnetBlock(
+                    ngf * mult,
+                    padding_type=padding_type,
+                    norm_layer=norm_layer,
+                    dropout_rate=dropout_rate,
+                    use_bias=use_bias)
+            ])
+
+        for i in range(n_blocks2):
+            self.model.extend([
+                SuperMobileResnetBlock(
+                    ngf * mult,
+                    padding_type=padding_type,
+                    norm_layer=norm_layer,
+                    dropout_rate=dropout_rate,
+                    use_bias=use_bias)
+            ])
+
+        for i in range(n_blocks3):
+            self.model.extend([
+                SuperMobileResnetBlock(
+                    ngf * mult,
+                    padding_type=padding_type,
+                    norm_layer=norm_layer,
+                    dropout_rate=dropout_rate,
+                    use_bias=use_bias)
+            ])
+
+        for i in range(n_downsampling):
+            mult = 2**(n_downsampling - i)
+            output_size = (i + 1) * 128
+            #### torch:out_padding = 1 => paddle:deconv + pad
+            self.model.extend([
+                SuperConv2DTranspose(
+                    ngf * mult,
+                    int(ngf * mult / 2),
+                    filter_size=3,
+                    output_size=output_size,
+                    stride=2,
+                    padding=1,
+                    bias_attr=use_bias), norm_layer(int(ngf * mult / 2)),
+                ReLU()
+            ])
+
+        self.model.extend([Pad2D(paddings=[3, 3, 3, 3], mode="reflect")])
+        self.model.extend(
+            [SuperConv2D(
+                ngf, output_nc, filter_size=7, padding=0)])
+
+    def forward(self, input):
+        configs = self.configs
+        x = tensor.clamp(input, min=-1, max=1)
+        cnt = 0
+        for i in range(0, 10):
+            sublayer = self.model[i]
+            if isinstance(sublayer, SuperConv2D):
+                channel = configs['channels'][cnt] * (2**cnt)
+                config = {'channel': channel}
+                x = sublayer(x, config)
+                cnt += 1
+            else:
+                x = sublayer(x)
+
+        for i in range(3):
+            for j in range(10 + i * 3, 13 + i * 3):
+                if len(configs['channels']) == 6:
+                    channel = configs['channels'][3] * 4
+                else:
+                    channel = configs['channels'][i + 3] * 4
+                config = {'channel': channel}
+                sublayer = self.model[j]
+                x = sublayer(x, config)
+
+        cnt = 2
+        for i in range(19, 27):
+            sublayer = self.model[i]
+            if isinstance(sublayer, SuperConv2DTranspose):
+                cnt -= 1
+                if len(configs['channels']) == 6:
+                    channel = configs['channels'][5 - cnt] * (2**cnt)
+                else:
+                    channel = configs['channels'][7 - cnt] * (2**cnt)
+                config = {'channel': channel}
+                x = sublayer(x, config)
+            elif isinstance(sublayer, SuperConv2D):
+                config = {'channel': sublayer._num_filters}
+                x = sublayer(x, config)
+            else:
+                x = sublayer(x)
+        x = fluid.layers.tanh(x)
+        return x

demo/gan_compression/models/layers.py

+# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import division
+import paddle.fluid as fluid
+import numpy as np
+from paddle.fluid.dygraph.nn import Conv2D, Conv2DTranspose, BatchNorm, InstanceNorm
+from paddle.nn.layer import Leaky_ReLU, ReLU, Pad2D
+import os
+import functools
+
+# cudnn is not better when batch size is 1.
+use_cudnn = False
+
+
+class conv2d(fluid.dygraph.Layer):
+    """docstring for Conv2D"""
+
+    def __init__(self,
+                 num_channels,
+                 num_filters=64,
+                 filter_size=7,
+                 stride=1,
+                 stddev=0.02,
+                 padding=0,
+                 norm=True,
+                 norm_layer=InstanceNorm,
+                 relu=True,
+                 relufactor=0.0,
+                 use_bias=False):
+        super(conv2d, self).__init__()
+
+        if use_bias == False:
+            con_bias_attr = False
+        else:
+            con_bias_attr = fluid.ParamAttr(
+                initializer=fluid.initializer.Constant(0.0))
+
+        self.conv = Conv2D(
+            num_channels=num_channels,
+            num_filters=int(num_filters),
+            filter_size=int(filter_size),
+            stride=stride,
+            padding=padding,
+            use_cudnn=use_cudnn,
+            param_attr=fluid.ParamAttr(
+                initializer=fluid.initializer.NormalInitializer(
+                    loc=0.0, scale=stddev)),
+            bias_attr=con_bias_attr)
+        if norm_layer == InstanceNorm:
+            self.bn = InstanceNorm(
+                num_channels=num_filters,
+                param_attr=fluid.ParamAttr(
+                    initializer=fluid.initializer.Constant(1.0),
+                    trainable=False),
+                bias_attr=fluid.ParamAttr(
+                    initializer=fluid.initializer.Constant(0.0),
+                    trainable=False), )
+        elif norm_layer == BatchNorm:
+            self.bn = BatchNorm(
+                num_channels=num_filters,
+                param_attr=fluid.ParamAttr(
+                    initializer=fluid.initializer.NormalInitializer(1.0,
+                                                                    0.02)),
+                bias_attr=fluid.ParamAttr(
+                    initializer=fluid.initializer.Constant(0.0)), )
+        else:
+            raise NotImplementedError
+
+        self.relufactor = relufactor
+        self.use_bias = use_bias
+        self.norm = norm
+        if relu:
+            if relufactor == 0.0:
+                self.lrelu = ReLU()
+            else:
+                self.lrelu = Leaky_ReLU(self.relufactor)
+        self.relu = relu
+
+    def forward(self, inputs):
+        conv = self.conv(inputs)
+        if self.norm:
+            conv = self.bn(conv)
+        if self.relu:
+            conv = self.lrelu(conv)
+            #conv = fluid.layers.leaky_relu(conv,alpha=self.relufactor)
+        return conv
+
+
+class SeparableConv2D(fluid.dygraph.Layer):
+    def __init__(self,
+                 num_channels,
+                 num_filters,
+                 filter_size,
+                 stride=1,
+                 padding=0,
+                 norm_layer='instance',
+                 use_bias=True,
+                 scale_factor=1,
+                 stddev=0.02):
+        super(SeparableConv2D, self).__init__()
+
+        if use_bias == False:
+            con_bias_attr = False
+        else:
+            con_bias_attr = fluid.ParamAttr(
+                initializer=fluid.initializer.Constant(0.0))
+
+        self.conv_sep = Conv2D(
+            num_channels=num_channels,
+            num_filters=num_channels * scale_factor,
+            filter_size=filter_size,
+            stride=stride,
+            padding=padding,
+            use_cudnn=use_cudnn,
+            groups=num_channels,
+            param_attr=fluid.ParamAttr(
+                initializer=fluid.initializer.NormalInitializer(
+                    loc=0.0, scale=stddev)),
+            bias_attr=con_bias_attr)
+
+        self.norm = InstanceNorm(
+            num_channels=num_filters,
+            param_attr=fluid.ParamAttr(
+                initializer=fluid.initializer.Constant(1.0)),
+            bias_attr=fluid.ParamAttr(
+                initializer=fluid.initializer.Constant(0.0)), )
+
+        self.conv_out = Conv2D(
+            num_channels=num_channels * scale_factor,
+            num_filters=num_filters,
+            filter_size=1,
+            stride=1,
+            use_cudnn=use_cudnn,
+            param_attr=fluid.ParamAttr(
+                initializer=fluid.initializer.NormalInitializer(
+                    loc=0.0, scale=stddev)),
+            bias_attr=con_bias_attr)
+
+    def forward(self, inputs):
+        conv = self.conv_sep(inputs)
+        conv = self.norm(conv)
+        conv = self.conv_out(conv)
+        return conv
+
+
+class DeConv2D(fluid.dygraph.Layer):
+    def __init__(self,
+                 num_channels,
+                 num_filters=64,
+                 filter_size=7,
+                 stride=1,
+                 stddev=0.02,
+                 padding=[0, 0],
+                 outpadding=[0, 0, 0, 0],
+                 relu=True,
+                 norm=True,
+                 norm_layer=InstanceNorm,
+                 relufactor=0.0,
+                 use_bias=False):
+        super(DeConv2D, self).__init__()
+
+        if use_bias == False:
+            de_bias_attr = False
+        else:
+            de_bias_attr = fluid.ParamAttr(
+                initializer=fluid.initializer.Constant(0.0))
+
+        self._deconv = Conv2DTranspose(
+            num_channels,
+            num_filters,
+            filter_size=filter_size,
+            stride=stride,
+            padding=padding,
+            param_attr=fluid.ParamAttr(
+                initializer=fluid.initializer.NormalInitializer(
+                    loc=0.0, scale=stddev)),
+            bias_attr=de_bias_attr)
+
+        self.pad = Pad2D(paddings=outpadding, mode='constant', pad_value=0.0)
+        if norm_layer == InstanceNorm:
+            self.bn = InstanceNorm(
+                num_channels=num_filters,
+                param_attr=fluid.ParamAttr(
+                    initializer=fluid.initializer.Constant(1.0),
+                    trainable=False),
+                bias_attr=fluid.ParamAttr(
+                    initializer=fluid.initializer.Constant(0.0),
+                    trainable=False), )
+        elif norm_layer == BatchNorm:
+            self.bn = BatchNorm(
+                num_channels=num_filters,
+                param_attr=fluid.ParamAttr(
+                    initializer=fluid.initializer.NormalInitializer(1.0,
+                                                                    0.02)),
+                bias_attr=fluid.ParamAttr(
+                    initializer=fluid.initializer.Constant(0.0)), )
+        else:
+            raise NotImplementedError
+
+        self.outpadding = outpadding
+        self.relufactor = relufactor
+        self.use_bias = use_bias
+        self.norm = norm
+        self.relu = relu
+        if relu:
+            if relufactor == 0.0:
+                self.lrelu = ReLU()
+            else:
+                self.lrelu = Leaky_ReLU(self.relufactor)
+
+    def forward(self, inputs):
+        #todo: add use_bias
+        #if self.use_bias==False:
+        conv = self._deconv(inputs)
+        #else:
+        #    conv = self._deconv(inputs)
+        #conv = fluid.layers.pad2d(conv, paddings=self.outpadding, mode='constant', pad_value=0.0)
+        conv = self.pad(conv)
+
+        if self.norm:
+            conv = self.bn(conv)
+        if self.relu:
+            #conv = fluid.layers.leaky_relu(conv,alpha=self.relufactor)
+            conv = self.lrelu(conv)
+        return conv

demo/gan_compression/models/loss.py

+import paddle.fluid as fluid
+
+
+def gan_loss(gan_mode, prediction, target_is_real, for_discriminator=True):
+    if target_is_real:
+        label = fluid.layers.fill_constant(
+            shape=fluid.layers.shape(prediction), value=1.0, dtype='float32')
+    else:
+        label = fluid.layers.fill_constant(
+            shape=fluid.layers.shape(prediction), value=0.0, dtype='float32')
+    if gan_mode == 'lsgan':
+        loss = fluid.layers.mse_loss(prediction, label)
+    elif gan_mode == 'vanilla':
+        loss = fluid.layers.sigmoid_cross_entropy_with_logits(prediction,
+                                                              label)
+    elif gan_mode == 'wgangp':
+        pass
+    elif gan_mode == 'hinge':
+        zero = fluid.layers.fill_constant(
+            shape=fluid.layers.shape(prediction), value=0.0, dtype='float32')
+        if for_discriminator:
+            if target_is_real:
+                minval = fluid.layers.elementwise_min(prediction - 1., zero)
+                loss = -1. * fluid.layers.reduce_mean(minval)
+            else:
+                minval = fluid.layers.elementwise_min(-1. * prediction - 1.,
+                                                      zero)
+                loss = -1. * fluid.layers.reduce_mean(minval)
+        else:
+            assert target_is_real
+            loss = -1. * fluid.layers.reduce_mean(prediction)
+    else:
+        raise NotImplementedError('gan mode %s not implemented' % gan_mode)
+    return loss
+
+
+def recon_loss(mode, prediction, label):
+    if mode == 'l1':
+        loss = fluid.layers.reduce_mean(
+            fluid.layers.elementwise_sub(
+                prediction, label, act='abs'))
+    elif mode == 'l2':
+        loss = fluid.layers.mse_loss(prediction, label)
+    elif mode == 'smooth_l1':
+        loss = fluid.layers.reduce_mean(
+            fluid.layers.smooth_l1(prediction, label))
+    elif mode == 'vgg':
+        pass
+    else:
+        raise NotImplementedError('Unknown reconstruction loss type [%s]!' %
+                                  mode)
+    return loss

demo/gan_compression/models/modules.py

+import paddle.fluid as fluid
+from paddle.fluid.dygraph.nn import Conv2D, Conv2DTranspose, BatchNorm, InstanceNorm, Dropout
+from paddle.nn.layer import Leaky_ReLU, ReLU, Pad2D
+
+__all__ = ['SeparableConv2D', 'MobileResnetBlock', 'ResnetBlock']
+
+use_cudnn = False
+
+
+class SeparableConv2D(fluid.dygraph.Layer):
+    def __init__(self,
+                 num_channels,
+                 num_filters,
+                 filter_size,
+                 stride=1,
+                 padding=0,
+                 norm_layer=InstanceNorm,
+                 use_bias=True,
+                 scale_factor=1,
+                 stddev=0.02,
+                 use_cudnn=use_cudnn):
+        super(SeparableConv2D, self).__init__()
+
+        self.conv = fluid.dygraph.LayerList([
+            Conv2D(
+                num_channels=num_channels,
+                num_filters=num_channels * scale_factor,
+                filter_size=filter_size,
+                stride=stride,
+                padding=padding,
+                use_cudnn=False,
+                groups=num_channels,
+                param_attr=fluid.ParamAttr(
+                    initializer=fluid.initializer.NormalInitializer(
+                        loc=0.0, scale=stddev)),
+                bias_attr=use_bias)
+        ])
+
+        self.conv.extend([norm_layer(num_channels * scale_factor)])
+
+        self.conv.extend([
+            Conv2D(
+                num_channels=num_channels * scale_factor,
+                num_filters=num_filters,
+                filter_size=1,
+                stride=1,
+                use_cudnn=use_cudnn,
+                param_attr=fluid.ParamAttr(
+                    initializer=fluid.initializer.NormalInitializer(
+                        loc=0.0, scale=stddev)),
+                bias_attr=use_bias)
+        ])
+
+    def forward(self, inputs):
+        for sublayer in self.conv:
+            inputs = sublayer(inputs)
+        return inputs
+
+
+class MobileResnetBlock(fluid.dygraph.Layer):
+    def __init__(self, in_c, out_c, padding_type, norm_layer, dropout_rate,
+                 use_bias):
+        super(MobileResnetBlock, self).__init__()
+        self.padding_type = padding_type
+        self.dropout_rate = dropout_rate
+        self.conv_block = fluid.dygraph.LayerList([])
+
+        p = 0
+        if self.padding_type == 'reflect':
+            self.conv_block.extend(
+                [Pad2D(
+                    paddings=[1, 1, 1, 1], mode='reflect')])
+        elif self.padding_type == 'replicate':
+            self.conv_block.extend(
+                [Pad2D(
+                    inputs, paddings=[1, 1, 1, 1], mode='edge')])
+        elif self.padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' %
+                                      self.padding_type)
+
+        self.conv_block.extend([
+            SeparableConv2D(
+                num_channels=in_c,
+                num_filters=out_c,
+                filter_size=3,
+                padding=p,
+                stride=1), norm_layer(out_c), ReLU()
+        ])
+
+        self.conv_block.extend([Dropout(p=self.dropout_rate)])
+
+        if self.padding_type == 'reflect':
+            self.conv_block.extend(
+                [Pad2D(
+                    paddings=[1, 1, 1, 1], mode='reflect')])
+        elif self.padding_type == 'replicate':
+            self.conv_block.extend(
+                [Pad2D(
+                    inputs, paddings=[1, 1, 1, 1], mode='edge')])
+        elif self.padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' %
+                                      self.padding_type)
+
+        self.conv_block.extend([
+            SeparableConv2D(
+                num_channels=out_c,
+                num_filters=in_c,
+                filter_size=3,
+                padding=p,
+                stride=1), norm_layer(in_c)
+        ])
+
+    def forward(self, inputs):
+        y = inputs
+        for sublayer in self.conv_block:
+            y = sublayer(y)
+        out = inputs + y
+        return out
+
+
+class ResnetBlock(fluid.dygraph.Layer):
+    def __init__(self,
+                 dim,
+                 padding_type,
+                 norm_layer,
+                 dropout_rate,
+                 use_bias=False):
+        super(ResnetBlock, self).__init__()
+
+        self.conv_block = fluid.dygraph.LayerList([])
+        p = 0
+        if padding_type == 'reflect':
+            self.conv_block.extend(
+                [Pad2D(
+                    paddings=[1, 1, 1, 1], mode='reflect')])
+        elif padding_type == 'replicate':
+            self.conv_block.extend([Pad2D(paddings=[1, 1, 1, 1], mode='edge')])
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' %
+                                      padding_type)
+
+        self.conv_block.extend([
+            Conv2D(
+                dim, dim, filter_size=3, padding=p, bias_attr=use_bias),
+            norm_layer(dim), ReLU()
+        ])
+        self.conv_block.extend([Dropout(dropout_rate)])
+
+        p = 0
+        if padding_type == 'reflect':
+            self.conv_block.extend(
+                [Pad2D(
+                    paddings=[1, 1, 1, 1], mode='reflect')])
+        elif padding_type == 'replicate':
+            self.conv_block.extend([Pad2D(paddings=[1, 1, 1, 1], mode='edge')])
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' %
+                                      padding_type)
+
+        self.conv_block.extend([
+            Conv2D(
+                dim, dim, filter_size=3, padding=p, bias_attr=use_bias),
+            norm_layer(dim)
+        ])
+
+    def forward(self, inputs):
+        y = inputs
+        for sublayer in self.conv_block:
+            y = sublayer(y)
+        return y + inputs

demo/gan_compression/models/network.py

+import functools
+import paddle.fluid as fluid
+from paddle.fluid.dygraph.nn import BatchNorm, InstanceNorm
+from discrimitor import NLayerDiscriminator
+from generator.resnet_generator import ResnetGenerator
+from generator.mobile_generator import MobileResnetGenerator
+from generator.super_generator import SuperMobileResnetGenerator
+
+
+class Identity(fluid.dygraph.Layer):
+    def forward(self, x):
+        return x
+
+
+def get_norm_layer(norm_type='instance'):
+    if norm_type == 'instance':
+        norm_layer = functools.partial(
+            InstanceNorm,
+            param_attr=fluid.ParamAttr(
+                initializer=fluid.initializer.Constant(1.0),
+                learning_rate=0.0,
+                trainable=False),
+            bias_attr=fluid.ParamAttr(
+                initializer=fluid.initializer.Constant(0.0),
+                learning_rate=0.0,
+                trainable=False))
+    elif norm_type == 'batch':
+        norm_layer = functools.partial(
+            BatchNorm,
+            param_attr=fluid.ParamAttr(
+                initializer=fluid.initializer.NormalInitializer(1.0, 0.02)),
+            bias_attr=fluid.ParamAttr(
+                initializer=fluid.initializer.Constant(0.0)))
+    elif norm_type == 'none':
+
+        def norm_layer(x):
+            return Identity(x)
+    else:
+        raise NotImplementedError('normalization layer [%s] is not found' %
+                                  norm_type)
+    return norm_layer
+
+
+def define_G(input_nc,
+             output_nc,
+             ngf,
+             netG,
+             norm_type='batch',
+             dropout_rate=0,
+             init_type='normal',
+             stddev=0.02):
+    net = None
+    norm_layer = get_norm_layer(norm_type)
+    if netG == 'resnet_9blocks':
+        net = ResnetGenerator(
+            input_nc,
+            output_nc,
+            ngf,
+            norm_layer=norm_layer,
+            dropout_rate=dropout_rate,
+            n_blocks=9)
+    elif netG == 'mobile_resnet_9blocks':
+        net = MobileResnetGenerator(
+            input_nc,
+            output_nc,
+            ngf,
+            norm_layer=norm_layer,
+            dropout_rate=dropout_rate,
+            n_blocks=9)
+    elif netG == 'super_mobile_resnet_9blocks':
+        net = SuperMobileResnetGenerator(
+            input_nc,
+            output_nc,
+            ngf,
+            norm_layer=norm_layer,
+            dropout_rate=dropout_rate,
+            n_blocks=9)
+    return net
+
+
+def define_D(input_nc, ndf, netD, norm_type='batch', n_layers_D=3):
+    net = None
+    norm_layer = get_norm_layer(norm_type)
+    if netD == 'n_layers':
+        net = NLayerDiscriminator(
+            input_nc, ndf, n_layers_D, norm_layer=norm_layer)
+    return net

demo/gan_compression/models/super_modules.py

+import paddle.fluid as fluid
+import paddle.fluid.dygraph_utils as dygraph_utils
+from paddle.fluid.data_feeder import check_variable_and_dtype, check_type
+from paddle.fluid.dygraph.base import to_variable
+from paddle.fluid.framework import in_dygraph_mode
+from paddle.fluid.dygraph.nn import InstanceNorm, Conv2D, Conv2DTranspose
+import paddle.fluid.core as core
+import numpy as np
+
+use_cudnn = False
+
+
+class SuperInstanceNorm(fluid.dygraph.InstanceNorm):
+    def __init__(self,
+                 num_channels,
+                 epsilon=1e-5,
+                 param_attr=None,
+                 bias_attr=None,
+                 dtype='float32'):
+        super(SuperInstanceNorm, self).__init__(
+            num_channels,
+            epsilon=1e-5,
+            param_attr=None,
+            bias_attr=None,
+            dtype='float32')
+
+    def forward(self, input):
+        in_nc = int(input.shape[1])
+        scale = self.scale[:in_nc]
+        bias = self.scale[:in_nc]
+        if in_dygraph_mode():
+            out, _, _ = core.ops.instance_norm(input, scale, bias, 'epsilon',
+                                               self._epsilon)
+            return out
+        check_variable_and_dtype(input, 'input', ['float32', 'float64'],
+                                 "SuperInstanceNorm")
+
+        attrs = {"epsilon": self._epsilon}
+
+        inputs = {"X": [input], "Scale": [scale], "Bias": [bias]}
+
+        saved_mean = self._helper.create_variable_for_type_inference(
+            dtype=self._dtype, stop_gradient=True)
+        saved_variance = self._helper.create_variable_for_type_inference(
+            dtype=self._dtype, stop_gradient=True)
+        instance_norm_out = self._helper.create_variable_for_type_inference(
+            self._dtype)
+
+        outputs = {
+            "Y": [instance_norm_out],
+            "SavedMean": [saved_mean],
+            "SavedVariance": [saved_variance]
+        }
+
+        self._helper.append_op(
+            type="instance_norm", inputs=inputs, outputs=outputs, attrs=attrs)
+        return instance_norm_out
+
+
+class SuperConv2D(fluid.dygraph.Conv2D):
+    def __init__(self,
+                 num_channels,
+                 num_filters,
+                 filter_size,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 groups=None,
+                 param_attr=None,
+                 bias_attr=None,
+                 use_cudnn=True,
+                 act=None,
+                 dtype='float32'):
+        super(SuperConv2D, self).__init__(
+            num_channels, num_filters, filter_size, stride, padding, dilation,
+            groups, param_attr, bias_attr, use_cudnn, act, dtype)
+
+    def forward(self, input, config):
+        in_nc = int(input.shape[1])
+        out_nc = config['channel']
+        weight = self.weight[:out_nc, :in_nc, :, :]
+        #print('super conv shape', weight.shape)
+        if in_dygraph_mode():
+            if self._l_type == 'conv2d':
+                attrs = ('strides', self._stride, 'paddings', self._padding,
+                         'dilations', self._dilation, 'groups', self._groups
+                         if self._groups else 1, 'use_cudnn', self._use_cudnn)
+                out = core.ops.conv2d(input, weight, *attrs)
+            elif self._l_type == 'depthwise_conv2d':
+                attrs = ('strides', self._stride, 'paddings', self._padding,
+                         'dilations', self._dilation, 'groups', self._groups,
+                         'use_cudnn', self._use_cudnn)
+                out = core.ops.depthwise_conv2d(input, weight, *attrs)
+            else:
+                raise ValueError("conv type error")
+
+            pre_bias = out
+            if self.bias is not None:
+                bias = self.bias[:out_nc]
+                pre_act = dygraph_utils._append_bias_in_dygraph(pre_bias, bias,
+                                                                1)
+            else:
+                pre_act = pre_bias
+
+            return dygraph_utils._append_activation_in_dygraph(pre_act,
+                                                               self._act)
+
+        inputs = {'Input': [input], 'Filter': [weight]}
+        attrs = {
+            'strides': self._stride,
+            'paddings': self._padding,
+            'dilations': self._dilation,
+            'groups': self._groups if self._groups else 1,
+            'use_cudnn': self._use_cudnn,
+            'use_mkldnn': False,
+        }
+        check_variable_and_dtype(
+            input, 'input', ['float16', 'float32', 'float64'], 'SuperConv2D')
+        pre_bias = self._helper.create_variable_for_type_inference(
+            dtype=self._dtype)
+
+        self._helper.append_op(
+            type=self._l_type,
+            inputs={
+                'Input': input,
+                'Filter': weight,
+            },
+            outputs={"Output": pre_bias},
+            attrs=attrs)
+
+        if self.bias is not None:
+            bias = self.bias[:out_nc]
+            pre_act = self._helper.create_variable_for_type_inference(
+                dtype=self._dtype)
+            self._helper.append_op(
+                type='elementwise_add',
+                inputs={'X': [pre_bias],
+                        'Y': [bias]},
+                outputs={'Out': [pre_act]},
+                attrs={'axis': 1})
+        else:
+            pre_act = pre_bias
+
+        # Currently, we don't support inplace in dygraph mode
+        return self._helper.append_activation(pre_act, act=self._act)
+
+
+class SuperConv2DTranspose(fluid.dygraph.Conv2DTranspose):
+    def __init__(self,
+                 num_channels,
+                 num_filters,
+                 filter_size,
+                 output_size=None,
+                 padding=0,
+                 stride=1,
+                 dilation=1,
+                 groups=None,
+                 param_attr=None,
+                 bias_attr=None,
+                 use_cudnn=True,
+                 act=None,
+                 dtype='float32'):
+        super(SuperConv2DTranspose,
+              self).__init__(num_channels, num_filters, filter_size,
+                             output_size, padding, stride, dilation, groups,
+                             param_attr, bias_attr, use_cudnn, act, dtype)
+
+    def forward(self, input, config):
+        in_nc = int(input.shape[1])
+        out_nc = int(config['channel'])
+        weight = self.weight[:in_nc, :out_nc, :, :]
+        if in_dygraph_mode():
+            op = getattr(core.ops, self._op_type)
+            out = op(input, weight, 'output_size', self._output_size,
+                     'strides', self._stride, 'paddings', self._padding,
+                     'dilations', self._dilation, 'groups', self._groups,
+                     'use_cudnn', self._use_cudnn)
+            pre_bias = out
+            if self.bias is not None:
+                bias = self.bias[:out_nc]
+                pre_act = dygraph_utils._append_bias_in_dygraph(pre_bias, bias,
+                                                                1)
+            else:
+                pre_act = pre_bias
+
+            return dygraph_utils._append_activation_in_dygraph(
+                pre_act, act=self._act)
+
+        check_variable_and_dtype(input, 'input',
+                                 ['float16', 'float32', 'float64'],
+                                 "SuperConv2DTranspose")
+
+        inputs = {'Input': [input], 'Filter': [weight]}
+        attrs = {
+            'output_size': self._output_size,
+            'strides': self._stride,
+            'paddings': self._padding,
+            'dilations': self._dilation,
+            'groups': self._groups,
+            'use_cudnn': self._use_cudnn
+        }
+
+        pre_bias = self._helper.create_variable_for_type_inference(
+            dtype=input.dtype)
+        self._helper.append_op(
+            type=self._op_type,
+            inputs=inputs,
+            outputs={'Output': pre_bias},
+            attrs=attrs)
+
+        if self.bias is not None:
+            pre_act = self._helper.create_variable_for_type_inference(
+                dtype=self._dtype)
+            self._helper.append_op(
+                type='elementwise_add',
+                inputs={'X': [pre_bias],
+                        'Y': [bias]},
+                outputs={'Out': [pre_act]},
+                attrs={'axis': 1})
+        else:
+            pre_act = pre_bias
+
+        out = self._helper.append_activation(pre_act, act=self._act)
+        return out
+
+
+class SuperSeparableConv2D(fluid.dygraph.Layer):
+    def __init__(self,
+                 num_channels,
+                 num_filters,
+                 filter_size,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 norm_layer=InstanceNorm,
+                 bias_attr=None,
+                 scale_factor=1,
+                 use_cudnn=False):
+        super(SuperSeparableConv2D, self).__init__()
+        self.conv = fluid.dygraph.LayerList([
+            fluid.dygraph.nn.Conv2D(
+                num_channels=num_channels,
+                num_filters=num_channels * scale_factor,
+                filter_size=filter_size,
+                stride=stride,
+                padding=padding,
+                use_cudnn=False,
+                groups=num_channels,
+                bias_attr=bias_attr)
+        ])
+        if norm_layer == InstanceNorm:
+            #self.conv.extend([SuperInstanceNorm(num_channels * scale_factor)])
+            self.conv.extend([
+                SuperInstanceNorm(
+                    num_channels * scale_factor,
+                    param_attr=fluid.ParamAttr(
+                        initializer=fluid.initializer.Constant(1.0),
+                        learning_rate=0.0,
+                        trainable=False),
+                    bias_attr=fluid.ParamAttr(
+                        initializer=fluid.initializer.Constant(0.0),
+                        learning_rate=0.0,
+                        trainable=False))
+            ])
+        else:
+            raise NotImplementedError
+        self.conv.extend([
+            Conv2D(
+                num_channels=num_channels * scale_factor,
+                num_filters=num_filters,
+                filter_size=1,
+                stride=1,
+                use_cudnn=use_cudnn,
+                bias_attr=bias_attr)
+        ])
+
+    def forward(self, input, config):
+        in_nc = int(input.shape[1])
+        out_nc = int(config['channel'])
+        weight = self.conv[0].weight[:in_nc]
+        ###  conv1
+        if in_dygraph_mode():
+            if self.conv[0]._l_type == 'conv2d':
+                attrs = ('strides', self.conv[0]._stride, 'paddings',
+                         self.conv[0]._padding, 'dilations',
+                         self.conv[0]._dilation, 'groups', in_nc, 'use_cudnn',
+                         self.conv[0]._use_cudnn)
+                out = core.ops.conv2d(input, weight, *attrs)
+            elif self.conv[0]._l_type == 'depthwise_conv2d':
+                attrs = ('strides', self.conv[0]._stride, 'paddings',
+                         self.conv[0]._padding, 'dilations',
+                         self.conv[0]._dilation, 'groups', in_nc, 'use_cudnn',
+                         self.conv[0]._use_cudnn)
+                out = core.ops.depthwise_conv2d(input, weight, *attrs)
+            else:
+                raise ValueError("conv type error")
+
+            pre_bias = out
+            if self.conv[0].bias is not None:
+                bias = self.conv[0].bias[:in_nc]
+                pre_act = dygraph_utils._append_bias_in_dygraph(pre_bias, bias,
+                                                                1)
+            else:
+                pre_act = pre_bias
+
+            conv0_out = dygraph_utils._append_activation_in_dygraph(
+                pre_act, self.conv[0]._act)
+
+        norm_out = self.conv[1](conv0_out)
+
+        weight = self.conv[2].weight[:out_nc, :in_nc, :, :]
+
+        if in_dygraph_mode():
+            if self.conv[2]._l_type == 'conv2d':
+                attrs = ('strides', self.conv[2]._stride, 'paddings',
+                         self.conv[2]._padding, 'dilations',
+                         self.conv[2]._dilation, 'groups', self.conv[2]._groups
+                         if self.conv[2]._groups else 1, 'use_cudnn',
+                         self.conv[2]._use_cudnn)
+                out = core.ops.conv2d(norm_out, weight, *attrs)
+            elif self.conv[2]._l_type == 'depthwise_conv2d':
+                attrs = ('strides', self.conv[2]._stride, 'paddings',
+                         self.conv[2]._padding, 'dilations',
+                         self.conv[2]._dilation, 'groups',
+                         self.conv[2]._groups, 'use_cudnn',
+                         self.conv[2]._use_cudnn)
+                out = core.ops.depthwise_conv2d(norm_out, weight, *attrs)
+            else:
+                raise ValueError("conv type error")
+
+            pre_bias = out
+            if self.conv[2].bias is not None:
+                bias = self.conv[2].bias[:out_nc]
+                pre_act = dygraph_utils._append_bias_in_dygraph(pre_bias, bias,
+                                                                1)
+            else:
+                pre_act = pre_bias
+
+            conv1_out = dygraph_utils._append_activation_in_dygraph(
+                pre_act, self.conv[2]._act)
+        return conv1_out
+
+
+if __name__ == '__main__':
+
+    class Net(fluid.dygraph.Layer):
+        def __init__(self, in_cn=3):
+            super(Net, self).__init__()
+            self.myconv = SuperSeparableConv2D(
+                num_channels=in_cn, num_filters=3, filter_size=3)
+
+        def forward(self, input, config):
+            print(input.shape[1])
+            conv = self.myconv(input, config)
+
+            return conv
+
+    config = {'channel': 2}
+    with fluid.dygraph.guard():
+        net = Net()
+
+        data_A = np.random.random((1, 3, 256, 256)).astype("float32")
+        data_A = to_variable(data_A)
+
+        out = net(data_A, config)
+        print(out.numpy())

demo/gan_compression/models/test_model.py

+from functools import reduce
+import paddle.fluid as fluid
+import network
+from utils import util
+#from utils.profile import profile_flops
+from paddleslim.analysis.flops import dygraph_flops
+
+
+class TestModel(fluid.dygraph.Layer):
+    def __init__(self, cfgs):
+        super(TestModel, self).__init__()
+        self.model_names = ['G']
+        self.netG = network.define_G(cfgs.input_nc, cfgs.output_nc, cfgs.ngf,
+                                     cfgs.netG, cfgs.norm_type,
+                                     cfgs.dropout_rate)
+        self.netG.eval()
+        self.cfgs = cfgs
+
+    def set_input(self, input):
+        self.real_A = input[0]
+
+    def setup(self):
+        self.load_network()
+
+    def load_network(self):
+        for name in self.model_names:
+            net = getattr(self, 'net' + name, None)
+            path = getattr(self.cfgs, 'restore_%s_path' % name, None)
+            if path is not None:
+                util.load_network(net, path)
+
+    def forward(self, config=None):
+        if config is not None:
+            self.netG.configs = config
+        self.fake_B = self.netG(self.real_A)
+
+    def test(self, config=None):
+        with fluid.dygraph.no_grad():
+            self.forward(config)
+
+    def profile(self, config=None):
+        netG = self.netG
+        netG.configs = config
+        with fluid.dygraph.no_grad():
+            flops = dygraph_flops(
+                netG, (self.real_A[:1]), only_conv=False, only_multiply=True)
+        params = 0
+        for p in netG.parameters():
+            if 'instance_norm' in p.name:
+                continue
+            params += reduce(lambda x, y: x * y, p.shape)
+        return flops, params

demo/gan_compression/select_arch.py

+import argparse
+import pickle
+
+
+def flops(item):
+    return item['flops']
+
+
+def main(cfgs):
+    with open(cfgs.pkl_path, 'rb') as f:
+        results = pickle.load(f)
+
+    result.sort(key=flops)
+
+    for item in results:
+        assert isinstance(item, dict)
+        qualified = True
+        if item['flops'] > cfgs.flops:
+            qualified = False
+        elif 'fid' in item and item['fid'] > cfgs.fid:
+            qualified = False
+        if qualified:
+            print(item)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        '--pkl_path', type=str, required=True, help='the input .pkl file path')
+    parser.add_argument(
+        '--flops', type=float, default=5.68e9, help='the FLOPs threshold')
+    parser.add_argument(
+        '--fid', type=float, default=-1, help='the FID threshold')
+    cfgs = parser.parse_args()
+    main(cfgs)

demo/gan_compression/supernets/__init__.py

+import importlib
+from models.base_model import BaseModel
+
+
+def find_model_using_name(model_name):
+    model_filename = "supernets." + model_name + "_supernet"
+    modellib = importlib.import_module(model_filename)
+    target_model_name = model_name.replace("_", "") + "supernet"
+    model = None
+    for name, cls in modellib.__dict__.items():
+        if name.lower() == target_model_name.lower() and issubclass(cls,
+                                                                    BaseModel):
+            model = cls
+    assert model is not None, "model {} is not right, please check it!".format(
+        model_name)
+
+    return model
+
+
+def get_special_cfg(model):
+    model_cls = find_model_using_name(model)
+    return model_cls.add_special_cfgs
+
+
+def create_supernet(cfg):
+    supernet = find_model_using_name(cfg.supernet)
+    return supernet(cfg)

demo/gan_compression/supernets/resnet_supernet.py

+import os
+import numpy as np
+import paddle.fluid as fluid
+from distillers.base_resnet_distiller import BaseResnetDistiller
+from models.super_modules import SuperConv2D
+from models import loss
+from configs.resnet_configs import get_configs
+from metric import get_fid
+from utils import util
+
+
+class ResnetSupernet(BaseResnetDistiller):
+    @staticmethod
+    def add_special_cfgs(parser, load_pre=False):
+        parser.add_argument(
+            '--supernet_lr',
+            type=float,
+            default=2e-4,
+            help="Initial learning rate to train super net")
+        parser.add_argument(
+            '--supernet_epoch',
+            type=int,
+            default=400,
+            help="The number of epoch to train super net")
+        parser.add_argument(
+            '--supernet_nepochs',
+            type=int,
+            default=200,
+            help="number of epochs with the initial learning rate")
+        parser.add_argument(
+            '--supernet_nepochs_decay',
+            type=int,
+            default=200,
+            help="number of epochs to linearly decay learning rate to zero")
+        parser.add_argument(
+            '--supernet_scheduler',
+            type=str,
+            default='linear',
+            help="learning rate scheduler in train supernet")
+        parser.add_argument(
+            '--supernet_student_netG',
+            type=str,
+            default='super_mobile_resnet_9blocks',
+            help="Which student generator network to choose in supernet")
+        parser.add_argument(
+            '--config_set',
+            type=str,
+            default='channels-32',
+            help="a set of configuration to get subnets of supernet")
+        parser.add_argument(
+            '--config_str',
+            type=str,
+            default=None,
+            help="the configuration string used to get specific subnet of supernet"
+        )
+        if load_pre:
+            super(ResnetSupernet, ResnetSupernet).add_special_cfgs(parser)
+        return parser
+
+    def __init__(self, cfgs):
+        assert 'super' in cfgs.supernet_student_netG
+        super(ResnetSupernet, self).__init__(cfgs, task='supernet')
+        self.best_fid_largest = 1e9
+        self.best_fid_smallest = 1e9
+        self.fids_largest, self.fids_smallest = [], []
+        if cfgs.config_set is not None:
+            assert cfgs.config_str is None
+            self.configs = get_configs(cfgs.config_set)
+            self.cfgs.eval_mode = 'both'
+        else:
+            assert cfgs.config_str is not None
+            self.configs = SingleConfigs(decode_config(cfgs.config_str))
+            self.opt.eval_mode = 'largest'
+
+    def forward(self, config):
+        with fluid.dygraph.no_grad():
+            self.Tfake_B = self.netG_teacher(self.real_A)
+        self.Tfake_B.stop_gradient = True
+        self.netG_student.configs = config
+        self.Sfake_B = self.netG_student(self.real_A)
+
+    def calc_distill_loss(self):
+        losses = []
+        for i, netA in enumerate(self.netAs):
+            assert isinstance(netA, SuperConv2D)
+            n = self.mapping_layers[i]
+            Tact = self.Tacts[n]
+            Sact = self.Sacts[n]
+            Sact = netA(Sact, {'channel': netA._num_filters})
+            loss = fluid.layers.mse_loss(Sact, Tact)
+            setattr(self, 'loss_G_distill%d' % i, loss)
+            losses.append(loss)
+        return sum(losses)
+
+    def backward_G(self):
+        self.loss_G_recon = loss.recon_loss(
+            self.cfgs.recon_loss_mode, self.Sfake_B,
+            self.Tfake_B) * self.cfgs.lambda_recon
+        pred_fake = self.netD(self.Sfake_B)
+        self.loss_G_gan = loss.gan_loss(
+            self.cfgs.gan_loss_mode, pred_fake, True,
+            for_discriminator=False) * self.cfgs.lambda_gan
+        if self.cfgs.lambda_distill > 0:
+            self.loss_G_distill = self.calc_distill_loss(
+            ) * self.cfgs.lambda_distill
+        else:
+            self.loss_G_distill = 0
+        self.loss_G = self.loss_G_gan + self.loss_G_recon + self.loss_G_distill
+        self.loss_G.backward()
+
+    def optimize_parameter(self):
+        config = self.configs.sample()
+        self.forward(config=config)
+        self.set_stop_gradient(self.netD, False)
+        self.backward_D()
+        self.set_stop_gradient(self.netD, True)
+        self.backward_G()
+
+        self.optimizer_D.optimizer.minimize(self.loss_D)
+        self.optimizer_D.optimizer.clear_gradients()
+        self.optimizer_G.optimizer.minimize(self.loss_G)
+        self.optimizer_G.optimizer.clear_gradients()
+
+    def evaluate_model(self, step):
+        ret = {}
+        self.is_best = False
+        save_dir = os.path.join(self.cfgs.save_dir, 'eval', str(step))
+        if not os.path.exists(save_dir):
+            os.makedirs(save_dir)
+        self.netG_student.eval()
+        if self.cfgs.eval_mode == 'both':
+            setting = ('largest', 'smallest')
+        else:
+            setting = (self.cfgs.eval_mode, )
+        for config_name in setting:
+            config = self.configs(config_name)
+            fakes, names = [], []
+            for i, data_i in enumerate(self.eval_dataloader):
+                id2name = self.name
+                self.set_single_input(data_i)
+                self.test(config)
+                fakes.append(self.Sfake_B.detach().numpy())
+                for j in range(len(self.Sfake_B)):
+                    if i < 10:
+                        Sname = 'Sfake_' + str(id2name[i + j]) + '.png'
+                        Tname = 'Tfake_' + str(id2name[i + j]) + '.png'
+                        Sfake_im = util.tensor2img(self.Sfake_B[j])
+                        Tfake_im = util.tensor2img(self.Tfake_B[j])
+                        util.save_image(Sfake_im,
+                                        os.path.join(save_dir, Sname))
+                        util.save_image(Tfake_im,
+                                        os.path.join(save_dir, Tname))
+
+            suffix = self.cfgs.direction
+            fluid.disable_imperative()
+            fid = get_fid(fakes, self.inception_model, self.npz,
+                          self.cfgs.inception_model)
+            fluid.enable_imperative()
+            if fid < getattr(self, 'best_fid_%s' % config_name, fid):
+                self.is_best = True
+                setattr(self, 'best_fid_%s' % config_name, fid)
+            fids = getattr(self, 'fids_%s' % config_name)
+            fids.append(fid)
+            if len(fids) > 3:
+                fids.pop(0)
+            ret['metric/fid_%s' % config_name] = fid
+            ret['metric/fid_%s-mean' % config_name] = sum(
+                getattr(self, 'fids_%s' % config_name)) / len(
+                    getattr(self, 'fids_%s' % config_name))
+            ret['metric/fid_%s-best' % config_name] = getattr(
+                self, 'best_fid_%s' % config_name)
+            print(
+                "SuperNet Evalution config_name is : %s, fid score is: %f, best fid score is %f"
+                %
+                (config_name, fid, getattr(self, 'best_fid_%s' % config_name)))
+
+        self.netG_student.train()
+        return ret
+
+    def test(self, config):
+        with fluid.dygraph.no_grad():
+            self.forward(config)

demo/gan_compression/utils/get_args.py

+import argparse
+import ast
+
+import models
+import distillers
+import supernets
+
+__all__ = ['configs']
+
+
+class configs:
+    def base_config(self, parser):
+        parser.add_argument(
+            '--model',
+            type=str,
+            default='cycle_gan',
+            help="The model want to compression")
+        parser.add_argument(
+            '--task',
+            type=str,
+            default='mobile+distiller+supernet',
+            help="Base channels in generator")
+        parser.add_argument(
+            '--distiller',
+            type=str,
+            default='resnet',
+            help="generator network in distiller")
+        parser.add_argument(
+            '--supernet',
+            type=str,
+            default='resnet',
+            help="generator network in supernet")
+        parser.add_argument(
+            '--use_gpu',
+            type=ast.literal_eval,
+            default=True,
+            help='Whether to use GPU in train/test model.')
+        ### data
+        parser.add_argument(
+            '--batch_size',
+            type=int,
+            default=1,
+            help="Minbatch size in all training")
+        parser.add_argument(
+            '--shuffle',
+            type=ast.literal_eval,
+            default=True,
+            help="Whether to shuffle data in training")
+        parser.add_argument(
+            '--flip',
+            type=ast.literal_eval,
+            default=True,
+            help="Whether to flip data randomly in training")
+        parser.add_argument(
+            '--dataset',
+            type=str,
+            default='horse2zebra',
+            help="The name of dataset")
+        parser.add_argument(
+            '--dataroot',
+            type=str,
+            default='./data',
+            help="The dictionary of data")
+        parser.add_argument(
+            '--image_size',
+            type=int,
+            default=286,
+            help="The image size when load image")
+        parser.add_argument(
+            '--crop_size',
+            type=int,
+            default=256,
+            help="The crop size used to crop image")
+        parser.add_argument(
+            '--crop_type',
+            type=str,
+            default='Random',
+            help="Which method to crop image")
+        ##############
+        parser.add_argument(
+            '--gan_loss_mode',
+            type=str,
+            default='lsgan',
+            help="The mode used to compute gan loss")
+        parser.add_argument(
+            '--ngf', type=int, default=64, help="Base channels in generator")
+        parser.add_argument(
+            '--netG',
+            type=str,
+            default='mobile_resnet_9blocks',
+            help="Which generator network to choose")
+        parser.add_argument(
+            '--dropout_rate',
+            type=float,
+            default=0,
+            help="dropout rate in generator")
+        ###############
+        parser.add_argument(
+            '--input_nc', type=int, default=3, help="Channel of input")
+        parser.add_argument(
+            '--output_nc', type=int, default=3, help="Channel of output")
+        parser.add_argument(
+            '--norm_type',
+            type=str,
+            default='instance',
+            help="The type of normalization")
+        parser.add_argument(
+            '--save_dir',
+            type=str,
+            default='./output',
+            help="The directory the model and the test result to saved")
+        parser.add_argument(
+            '--netD',
+            type=str,
+            default='n_layers',
+            help="Which discriminator network to choose")
+        parser.add_argument(
+            '--ndf',
+            type=int,
+            default=64,
+            help="Base channels in discriminator")
+        parser.add_argument(
+            '--n_layer_D',
+            type=int,
+            default=3,
+            help="The number of layer in discriminator, only used when netD == n_layers"
+        )
+        parser.add_argument(
+            '--beta1', type=float, default=0.5, help="momentum term of adam")
+        parser.add_argument(
+            '--direction',
+            type=str,
+            default='AtoB',
+            help="the direction of generator")
+        parser.add_argument(
+            '--step_per_epoch',
+            type=int,
+            default=1333,
+            help="The number of step in each epoch")
+        parser.add_argument(
+            '--inception_model',
+            type=str,
+            default='metric/params_inceptionV3',
+            help="The directory of inception model, used in computing fid")
+
+        parser.add_argument(
+            '--restore_D_path',
+            type=str,
+            default=None,
+            help="the pretrain model path of discriminator")
+        parser.add_argument(
+            '--restore_G_path',
+            type=str,
+            default=None,
+            help="the pretrain model path of generator")
+        parser.add_argument(
+            '--restore_A_path',
+            type=str,
+            default=None,
+            help="the pretrain model path of list of conv used in distiller")
+        parser.add_argument(
+            '--restore_O_path',
+            type=str,
+            default=None,
+            help="the pretrain model path of optimization")
+
+        parser.add_argument(
+            '--print_freq', type=int, default=1, help="print log frequency")
+        parser.add_argument(
+            '--save_freq',
+            type=int,
+            default=1,
+            help="the epoch frequency to save model")
+
+        return parser
+
+    def get_all_config(self):
+        parser = argparse.ArgumentParser(description="Configs for all model")
+        parser = self.base_config(parser)
+        cfg, _ = parser.parse_known_args()
+
+        task = cfg.task
+        tasks = task.split('+')
+        load_resbase = True
+        for t in tasks:
+            if t == 'mobile':
+                model = cfg.model
+                model_parser = models.get_special_cfg(model)
+                parser = model_parser(parser)
+            elif t == 'distiller':
+                model = cfg.distiller
+                model_parser = distillers.get_special_cfg(model)
+                parser = model_parser(parser, load_resbase)
+                load_resbase = False
+            elif t == 'supernet':
+                model = cfg.supernet
+                model_parser = supernets.get_special_cfg(model)
+                parser = model_parser(parser, load_resbase)
+                load_resbase = False
+            else:
+                raise NotImplementedError(
+                    "task name {} is error, please check it".format(t))
+        self.parser = parser
+        return parser.parse_args()
+
+    def print_configs(self, cfgs):
+        print("-----------  Configuration Arguments -----------")
+        for arg, value in sorted(vars(cfgs).items()):
+            print("%s: %s" % (arg, value))
+        print("------------------------------------------------")

demo/gan_compression/utils/image_pool.py

+import random
+import paddle.fluid as fluid
+
+
+class ImagePool():
+    def __init__(self, pool_size):
+        self.pool_size = pool_size
+        if self.pool_size > 0:
+            self.num_imgs = 0
+            self.images = []
+
+    def query(self, image):
+        if self.pool_size == 0:
+            return image
+        if self.num_imgs < self.pool_size:
+            self.num_imgs = self.num_imgs + 1
+            self.images.append(image)
+            return image
+        else:
+            p = random.uniform(0, 1)
+            if p > 0.5:
+                random_id = random.randint(0, self.pool_size - 1)
+                temp = self.images[random_id]
+                self.images[random_id] = image
+                return temp
+            else:
+                return image
+
+
+#        return_images = []
+#        for img in image:
+#            img = img.detach()
+#            img = fluid.layers.unsqueeze(img, axes=0)
+#            if self.num_imgs < self.pool_size:
+#                self.num_imgs = self.num_imgs + 1
+#                self.images.append(img)
+#                return_images.append(img)
+#            else:
+#                p = random.uniform(0, 1)
+#                if p > 0.5:
+#                    random_id = random.randint(0, self.pool_size - 1)
+#                    temp = self.images[random_id]
+#                    self.images[random_id] = img
+#                    return_images.append(temp)
+#                else:
+#                    return_images.append(img)
+#
+#        return_images = fluid.layers.concat(return_images, axis=0)
+#        return return_images

demo/gan_compression/utils/optimization.py

+import numpy as np
+import paddle.fluid as fluid
+from paddle.fluid.dygraph.learning_rate_scheduler import LearningRateDecay
+
+#step_per_epoch = 1334
+
+
+class LinearDecay(LearningRateDecay):
+    def __init__(self, learning_rate, step_per_epoch, nepochs, nepochs_decay):
+        super(LinearDecay, self).__init__()
+        self.learning_rate = learning_rate
+        self.nepochs = nepochs
+        self.nepochs_decay = nepochs_decay
+        self.step_per_epoch = step_per_epoch
+
+    def step(self):
+        cur_epoch = np.floor(self.step_num / self.step_per_epoch)
+        lr_l = 1.0 - max(0, cur_epoch + 1 -
+                         self.nepochs) / float(self.nepochs_decay + 1)
+        return self.create_lr_var(lr_l * self.learning_rate)
+
+
+class Optimizer:
+    def __init__(self,
+                 lr,
+                 scheduler,
+                 step_per_epoch,
+                 nepochs,
+                 nepochs_decay,
+                 args,
+                 parameter_list=None):
+        self.lr = lr
+        self.scheduler = scheduler
+        self.step_per_epoch = step_per_epoch
+        self.nepochs = nepochs
+        self.nepochs_decay = nepochs_decay
+        self.args = args
+        self.parameter_list = parameter_list
+        self.optimizer = self.lr_scheduler()
+
+    ### NOTE(ceci3): add more scheduler
+    def lr_scheduler(self):
+        if self.scheduler == 'linear':
+            self.scheduler_lr = LinearDecay(self.lr, self.step_per_epoch,
+                                            self.nepochs, self.nepochs_decay)
+        elif self.scheduler == 'step':
+            pass
+        elif self.scheduler == 'cosine':
+            pass
+        else:
+            return NotImplementedError(
+                'learning rate policy [%s] is not implemented', opt.lr_policy)
+
+        optimizer = fluid.optimizer.Adam(
+            learning_rate=self.scheduler_lr,
+            beta1=self.args.beta1,
+            beta2=0.999,
+            parameter_list=self.parameter_list)
+        return optimizer

demo/gan_compression/utils/util.py

+import os
+import numpy as np
+import pickle
+from PIL import Image
+import paddle.fluid as fluid
+
+
+def load_network(model, model_path):
+    if model_path.split('.')[-1] == 'pkl' or model_path.split('.')[
+            -1] == 'pth':
+        model_weight = pickle.load(open(model_path, 'rb'))
+        for key, value in model_weight.items():
+            model_weight[key] = np.array(value)
+    else:
+        assert os.path.exists(
+            model_path + '.pdparams'), "model path: {} is not exist!!!".format(
+                model_path + '.pdparams')
+        model_weight, _ = fluid.load_dygraph(model_path)
+    model.set_dict(model_weight)
+    print("params {} load done".format(model_path))
+    return model
+
+
+def load_optimizer(optimizer, optimizer_path):
+    assert os.path.exists(
+        optimizer + '.pdopt'), "optimizer path: {} is not exist!!!".format(
+            optimizer_path + '.pdopt')
+    _, optimier_info = fluid.load_dygraph(optimizer_path)
+    optimizer.set_dict(optimizer_info)
+    return optimizer
+
+
+def save_image(image, image_path):
+    if len(image.shape) == 4:
+        image = image[0]
+    if len(image.shape) == 2:
+        image = np.expand_dims(image, axis=2)
+    if image.shape[2] == 1:
+        image = np.repeat(image, 3, 2)
+    image_pil = Image.fromarray(image)
+    image_pil.save(image_path)
+
+
+def tensor2img(image_tensor, imtype=np.uint8, normalize=True, tile=False):
+    if isinstance(image_tensor, list):
+        image_numpy = []
+        for i in range(len(image_tensor)):
+            image_numpy.append(tensor2img(image_tensor[i], imtype, normalize))
+        return image_numpy
+
+    if len(image_tensor.shape) == 4:
+        images_np = []
+        for b in range(image_tensor.shape[0]):
+            one_image = image_tensor[b]
+            one_image_np = tensor2img(one_image)
+            images_np.append(np.expand_dims(one_image_np, axis=0))
+        images_np = np.concatenate(images_np, axis=0)
+        if tile:
+            images_tiled = tile_images(images_np)
+            return images_tiled
+        else:
+            return images_np
+
+    if len(image_tensor.shape) == 2:
+        #image_tensor = fluid.layers.unsqueeze(image_tensor, axes=0)
+        image_tensor = np.expand_dims(image_tensor, axis=0)
+    if type(image_tensor) != np.ndarray:
+        image_np = image_tensor.numpy()
+    else:
+        image_np = image_tensor
+    if normalize:
+        np.transpose(image_np, (1, 2, 0))
+        image_np = (np.transpose(image_np, (1, 2, 0)) + 1) / 2.0 * 255.0
+    else:
+        image_np = np.transpose(image_np, (1, 2, 0)) * 255.0
+    image_np = np.clip(image_np, 0, 255)
+    if image_np.shape[2] == 1:
+        image_np = image_np[:, :, 0]
+    return image_np.astype(imtype)

demo/gan_compression/utils/weight_transfer.py

+import paddle.fluid as fluid
+from paddle.fluid.dygraph.nn import Conv2D, Conv2DTranspose, InstanceNorm
+from models.modules import SeparableConv2D, MobileResnetBlock, ResnetBlock
+from paddle.fluid.dygraph.base import to_variable
+import numpy as np
+
+
+### CoutCinKhKw
+def transfer_Conv2D(m1, m2, input_index=None, output_index=None):
+    assert isinstance(m1, Conv2D) and isinstance(m2, Conv2D)
+    if m1.parameters()[0].shape[0] == 3:  ### last convolution
+        assert input_index is not None
+        m2.parameters()[0].set_value(m1.parameters()[0].numpy()[:,
+                                                                input_index])
+        if len(m2.parameters()) == 2:
+            m2.parameters()[1].set_value(m1.parameters()[1].numpy())
+        return None
+
+    else:
+        if m1.parameters()[0].shape[1] == 3:  ### first convolution
+            assert input_index is None
+            input_index = [0, 1, 2]
+        p = m1.parameters()[0]
+
+        if input_index is None:
+            q = fluid.layers.reduce_sum(fluid.layers.abs(p), dim=[0, 2, 3])
+            _, idx = fluid.layers.topk(q, m2.parameters()[0].shape[1])
+            p = p.numpy()[:, idx.numpy()]
+        else:
+            p = p.numpy()[:, input_index]
+
+        if output_index is None:
+            q = fluid.layers.reduce_sum(
+                fluid.layers.abs(to_variable(p)), dim=[1, 2, 3])
+            _, idx = fluid.layers.topk(q, m2.parameters()[0].shape[0])
+            idx = idx.numpy()
+        else:
+            idx = output_index
+
+        m2.parameters()[0].set_value(p[idx])
+        if len(m2.parameters()) == 2:
+            m2.parameters()[1].set_value(m1.parameters()[1].numpy()[idx])
+
+        return idx
+
+
+### CinCoutKhKw
+def transfer_Conv2DTranspose(m1, m2, input_index=None, output_index=None):
+    assert isinstance(m1, Conv2DTranspose) and isinstance(m2, Conv2DTranspose)
+    assert output_index is None
+    p = m1.parameters()[0]
+    with fluid.dygraph.guard():
+        if input_index is None:
+            q = fluid.layers.reduce_sum(fluid.layers.abs(p), dim=[1, 2, 3])
+            _, idx = fluid.layers.topk(q, m2.parameters()[0].shape[0])  ### Cin
+            p = p.numpy()[idx.numpy()]
+        else:
+            p = p.numpy()[input_index]
+
+        q = fluid.layers.reduce_sum(
+            fluid.layers.abs(to_variable(p)), dim=[0, 2, 3])
+        _, idx = fluid.layers.topk(q, m2.parameters()[0].shape[1])
+        idx = idx.numpy()
+        m2.parameters()[0].set_value(p[:, idx])
+        if len(m2.parameters()) == 2:
+            m2.parameters()[1].set_value(m1.parameters()[1].numpy()[idx])
+
+        return idx
+
+
+def transfer_SeparableConv2D(m1, m2, input_index=None, output_index=None):
+    assert isinstance(m1, SeparableConv2D) and isinstance(m2, SeparableConv2D)
+    dw1, pw1 = m1.conv[0], m1.conv[2]
+    dw2, pw2 = m2.conv[0], m2.conv[2]
+
+    if input_index is None:
+        p = dw1.parameters()[0]
+        q = fluid.layers.reduce_sum(fluid.layers.abs(p), dim=[1, 2, 3])
+        _, idx = fluid.layers.topk(q, dw2.parameters()[0].shape[0])
+        input_index = idx.numpy()
+    dw2.parameters()[0].set_value(dw1.parameters()[0].numpy()[input_index])
+
+    if len(dw2.parameters()) == 2:
+        dw2.parameters()[1].set_value(dw1.parameters()[1].numpy()[input_index])
+
+    idx = transfer_Conv2D(pw1, pw2, input_index, output_index)
+    return idx
+
+
+def transfer_MobileResnetBlock(m1, m2, input_index=None, output_index=None):
+    assert isinstance(m1, MobileResnetBlock) and isinstance(m2,
+                                                            MobileResnetBlock)
+    assert output_index is None
+    idx = transfer_SeparableConv2D(
+        m1.conv_block[1], m2.conv_block[1], input_index=input_index)
+    idx = transfer_SeparableConv2D(
+        m1.conv_block[6],
+        m2.conv_block[6],
+        input_index=idx,
+        output_index=input_index)
+    return idx
+
+
+def transfer_ResnetBlock(m1, m2, input_index=None, output_index=None):
+    assert isinstance(m1, ResnetBlock) and isinstance(m2, ResnetBlock)
+    assert output_index is None
+    idx = transfer_Conv2D(
+        m1.conv_block[1], m2.conv_block[1], input_index=input_index)
+    idx = transfer_Conv2D(
+        m1.conv_block[6],
+        m2.conv_block[6],
+        input_index=idx,
+        output_index=input_index)
+    return idx
+
+
+def transfer(m1, m2, input_index=None, output_index=None):
+    assert type(m1) == type(m2)
+    if isinstance(m1, Conv2D):
+        return transfer_Conv2D(m1, m2, input_index, output_index)
+    elif isinstance(m1, Conv2DTranspose):
+        return transfer_Conv2DTranspose(m1, m2, input_index, output_index)
+    elif isinstance(m1, ResnetBlock):
+        return transfer_ResnnetBlock(m1, m2, input_index, output_index)
+    elif isinstance(m1, MobileResnetBlock):
+        return transfer_MobileResnetBlock(m1, m2, input_index, output_index)
+    else:
+        raise NotImplementedError('Unknown module [%s]!' % type(m1))
+
+
+def load_pretrained_weight(model1, model2, netA, netB, ngf1, ngf2):
+    assert model1 == model2
+    assert ngf1 >= ngf2
+
+    index = None
+    if model1 == 'mobile_resnet_9blocks':
+        assert len(netA.sublayers()) == len(netB.sublayers())
+        for (n1, m1), (n2, m2) in zip(netA.named_sublayers(),
+                                      netB.named_sublayers()):
+            assert type(m1) == type(m2)
+            if len(n1) > 8:
+                continue
+            if isinstance(m1, (Conv2D, Conv2DTranspose, MobileResnetBlock)):
+                index = transfer(m1, m2, index)
+
+    elif model1 == 'resnet_9blocks':
+        assert len(netA.sublayers()) == len(netB.sublayers())
+        for m1, m2 in zip(netA.sublayers(), netB.sublayers()):
+            assert type(m1) == type(m2)
+            if isinstance(m1, (Conv2D, Conv2DTranspose, ResnetBlock)):
+                index = transfer(m1, m2, index)
+    else:
+        raise NotImplementedError('Unknown model [%s]!' % model1)
+
+
+class Test(fluid.dygraph.Layer):
+    def __init__(self):
+        super(Test, self).__init__()
+        self.net1 = SeparableConv2D(
+            num_channels=4, num_filters=6, filter_size=3)
+
+    def forward(self, x):
+        out = self.net1(x)
+        return out
+
+
+if __name__ == '__main__':
+    data = np.random.random((1, 4, 6, 6)).astype('float32')
+    with fluid.dygraph.guard():
+        net1 = Test()
+        net2 = Test()
+        net_1 = net1(to_variable(data))
+        net_2 = net2(to_variable(data))
+        model1 = model2 = 'mobile_resnet_9blocks'
+        load_pretrained_weight(model1, model2, net1, net2, 4, 4)