compressor.py

#   Copyright (c) 2022 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 logging
import os
import sys
import numpy as np
import inspect
from collections import namedtuple, Iterable
import paddle
import paddle.distributed.fleet as fleet
from ..quant.quant_post_hpo import quant_post_hpo
from ..quant.quanter import convert
from ..common.recover_program import recover_inference_program
from ..common import get_logger
from .create_compressed_program import build_distill_program, build_quant_program, build_prune_program
from .strategy_config import ProgramInfo, merge_config

_logger = get_logger(__name__, level=logging.INFO)


class AutoCompression:
    def __init__(self,
                 model_dir,
                 model_filename,
                 params_filename,
                 save_dir,
                 strategy_config,
                 train_config,
                 train_dataloader,
                 eval_callback,
                 devices='gpu'):
        ### model_dir(str): 模型路径
        ### model_filename(str): 模型文件名称
        ### params_filename(str): 参数文件名称
        ### save_dir(str): 压缩后模型保存的路径
        ### strategy_config(dict[dict]): 压缩策略配置, 包括量化配置、蒸馏配置
        ### train_config(dict): 训练配置
        ### train_dataloader(paddle.nn.Dataloader): 训练数据dataloader
        ### eval_callback(function，paddle.nn.Dataloader): eval回调函数，和测试数据之间必须传入一个，如果传入回调函数，则使用回调函数判断模型训练情况。callback传入predict结果（paddle的tensor），默认：None。
        self.model_dir = model_dir
        self.model_filename = model_filename
        self.params_filename = params_filename
        self.save_dir = save_dir
        self.strategy_config = strategy_config
        self.train_config = train_config
        self.train_dataloader = train_dataloader
        paddle.enable_static()
        if self.train_config is not None and self.train_config.use_fleet:
            fleet.init(is_collective=True)
        if self._prepare_eval(eval_callback) == 'eval_dataloader':
            self.eval_function = None
            self.eval_dataloader = eval_callback
        else:
            self.eval_function = eval_callback
            self.eval_dataloader = None

        self._strategy, self._config = self._prepare_strategy()
        self._exe, self._places = self._prepare_envs(devices)

    def _prepare_envs(self, devices):
        places = paddle.device._convert_to_place(devices)
        exe = paddle.static.Executor(places)
        return exe, places

    def _prepare_strategy(self):
        quant_config = self.strategy_config.get("Quantization", None)
        hpo_config = self.strategy_config.get("HyperParameterOptimization",
                                              None)
        prune_config = self.strategy_config.get("Prune", None)
        unstructure_prune_config = self.strategy_config.get("UnstructurePrune",
                                                            None)
        single_teacher_distill_config = self.strategy_config.get("Distillation",
                                                                 None)
        multi_teacher_distill_config = self.strategy_config.get(
            "MultiTeacherDistillation", None)

        assert (single_teacher_distill_config is None) or (multi_teacher_distill_config is None), \
            "Distillation and MultiTeacherDistillation cannot be set at the same time."
        self._distill_config = single_teacher_distill_config if \
               single_teacher_distill_config is not None else \
               multi_teacher_distill_config

        ### case1: quant_config & hpo_config ==> PTQ & HPO
        if quant_config is not None and hpo_config is not None:
            strategy = 'ptq_hpo'
            config = merge_config(quant_config, hpo_config)

        ### case2: quant_config & distill config ==> QAT & Distill
        elif quant_config is not None and self._distill_config is not None:
            strategy = 'qat_dis'
            config = merge_config(quant_config, self._distill_config)

        ### case3: prune_config & distill config
        elif prune_config is not None and self._distill_config is not None:
            strategy = 'prune_dis'
            config = merge_config(prune_config, self._distill_config)

        ### case4: unstructure_config & distill config
        elif unstructure_prune_config is not None and self._distill_config is not None:
            strategy = 'unstructure_prune_dis'
            config = merge_config(unstructure_prune_config,
                                  self._distill_config)

        ### case4: distill_config
        elif self._distill_config is not None:
            if single_teacher_distill_config is not None:
                strategy = 'single_teacher_dis'
                config = single_teacher_distill_config
            else:
                strategy = 'multi_teacher_dis'
                config = multi_teacher_distill_config

        ### case N: todo
        else:
            raise NotImplementedError(
                "Not Implemented {} be set at the same time now".format(
                    self.strategy_config.keys()))

        return strategy, config

    def _prepare_fleet_strategy(train_config):
        build_strategy = paddle.static.BuildStrategy()
        exec_strategy = paddle.static.ExecutionStrategy()

        strategy = fleet.DistributedStrategy()
        strategy.build_strategy = build_strategy
        if train_config.recompute_config is not None:
            strategy.recompute = True
            strategy.recompute_configs = { ** train_config.recompute_config}
        if train_config.sharding_config is not None:
            strategy.sharding = True
            strategy.sharding_configs = { ** train_config.sharding_config}
        if train_config.amp_config is not None:
            strategy.amp = True
            strategy.amp_configs = { ** train_config.amp_config}
        return strategy

    def _prepare_program(self, program, feed_target_names, fetch_targets):
        train_program = recover_inference_program(program)
        startup_program = paddle.static.Program()
        train_program_info = ProgramInfo(startup_program, train_program,
                                         feed_target_names, fetch_targets)

        config_dict = dict(self._config._asdict())
        ### add prune program
        self._pruner = None
        if 'prune' in self._strategy:
            self._pruner, train_program_info = build_prune_program(
                self._exe, self._places, config_dict, train_program_info,
                self._strategy)

        if self.train_config.use_fleet:
            dist_strategy = _prepare_fleet_strategy(self.train_config)
        else:
            dist_strategy = None

        ### add distill program
        if 'dis' in self._strategy:
            train_program_info, test_program_info = build_distill_program(
                self._exe,
                self._places,
                config_dict,
                self.train_config._asdict(),
                train_program_info,
                pruner=self._pruner,
                dist_strategy=dist_strategy)

        self._quant_config = None
        ### add quant_aware program, quant always is last step
        if 'qat' in self._strategy:
            train_program_info, test_program_info, self._quant_config = build_quant_program(
                self._exe, self._places, config_dict, train_program_info,
                test_program_info)

        self._exe.run(train_program_info.startup_program)

        if (not self.train_config.use_fleet
            ) and self.train_config.amp_config is not None:
            if hasattr(self.train_config.amp_config, 'use_pure_fp16'
                       ) and self.train_config.amp_config.use_pure_fp16:
                train_program_info.optimizer.amp_init(
                    self._places, scope=paddle.static.global_scope())

        if 'prune_algo' in config_dict and config_dict['prune_algo'] == 'asp':
            ### prune weight in scope
            self._pruner.prune_model(train_program_info.program)

        if not self.train_config.use_fleet:
            train_program_info = self._compiled_program(train_program_info,
                                                        self._strategy)
            test_program_info = self._compiled_program(test_program_info,
                                                       self._strategy)
        return train_program_info, test_program_info

    def _prepare_eval(self, eval_callback):
        if isinstance(eval_callback,
                      Iterable) or inspect.isgeneratorfunction(eval_callback):
            return 'eval_dataloader'
        else:
            return 'eval_callback'

    def _compiled_program(self, program_info, strategy):
        compiled_prog = paddle.static.CompiledProgram(program_info.program)
        build_strategy = paddle.static.BuildStrategy()
        exec_strategy = paddle.static.ExecutionStrategy()
        if 'qat' in strategy:
            build_strategy.memory_optimize = False
            build_strategy.enable_inplace = False
            build_strategy.fuse_all_reduce_ops = False
            build_strategy.sync_batch_norm = False

        compiled_prog = compiled_prog.with_data_parallel(
            loss_name=program_info.fetch_targets[0].name,
            build_strategy=build_strategy,
            exec_strategy=exec_strategy)
        program_info.program = compiled_prog
        return program_info

    def compress(self):
        ### start compress, including train/eval model
        if self._strategy == 'ptq_hpo':
            quant_post_hpo(
                self._exe,
                self._places,
                model_dir=self.model_dir,
                quantize_model_path=self.save_dir,
                train_dataloader=self.train_dataloader,
                eval_dataloader=self.eval_dataloader,
                eval_function=self.eval_function,
                model_filename=self.model_filename,
                params_filename=self.params_filename,
                save_model_filename=self.model_filename,
                save_params_filename=self.params_filename,
                quantizable_op_type=self._config.quantize_op_types,
                weight_bits=self._config.weight_bits,
                activation_bits=self._config.activation_bits,
                weight_quantize_type=self._config.weight_quantize_type,
                is_full_quantize=self._config.is_full_quantize,
                algo=self._config.ptq_algo,
                bias_correct=self._config.bias_correct,
                hist_percent=self._config.hist_percent,
                batch_size=[1],
                batch_num=self._config.batch_num,
                runcount_limit=self._config.max_quant_count)

        else:
            assert 'dis' in self._strategy, "Only support optimizer compressed model by distillation loss."

            ### convert a inference program to train program
            ###[inference_program, feed_target_names, fetch_targets]= paddle.static.load_inference_model( \
            ###    path_prefix=self.model_dir, \
            ###    model_filename=self.model_filename, params_filename=self.params_filename,
            ###    executor=self._exe)
            [inference_program, feed_target_names, fetch_targets]= paddle.fluid.io.load_inference_model( \
                dirname=self.model_dir, \
                model_filename=self.model_filename, params_filename=self.params_filename,
                executor=self._exe)

            ### used to check whether the dataloader is right
            if self.eval_function is not None and self.train_config.origin_metric is not None:
                metric = self.eval_function(self._exe, inference_program,
                                            feed_target_names, fetch_targets)
                _logger.info("metric of compressed model is: {}".format(metric))
                buf = 0.05
                if metric < (float(self.train_config.origin_metric) - buf) or \
                        metric > (float(self.train_config.origin_metric) + buf):
                    raise RuntimeError("target metric of pretrained model is {}, \
                          but now is {}, Please check the format of evaluation dataset \
                          or check the origin_metric in train_config"
                                                                     .format(\
                          self.train_config.origin_metric, metric))

            train_program_info, test_program_info = self._prepare_program(
                inference_program, feed_target_names, fetch_targets)

            test_program_info = self._start_train(train_program_info,
                                                  test_program_info)
            self._save_model(test_program_info)

    def _start_train(self, train_program_info, test_program_info):
        best_metric = -1.0
        for epoch_id in range(self.train_config.epochs):
            for batch_id, data in enumerate(self.train_dataloader()):
                np_probs_float, = self._exe.run(train_program_info.program, \
                    feed=data, \
                    fetch_list=train_program_info.fetch_targets)

                if 'unstructure' in self._strategy:
                    self._pruner.step()

                if self.train_config.logging_iter is None:
                    logging_iter = 10
                else:
                    logging_iter = self.train_config.logging_iter
                if batch_id % int(logging_iter) == 0:
                    _logger.info("epoch: {}, batch: {}, loss: {}".format(
                        epoch_id, batch_id, np_probs_float))

                if batch_id % int(self.train_config.eval_iter) == 0:
                    if self.eval_function is not None:

                        # GMP pruner step 3: update params before summrizing sparsity, saving model or evaluation. 
                        if 'unstructure' in self._strategy:
                            self._pruner.update_params()

                        metric = self.eval_function(
                            self._exe, test_program_info.program,
                            test_program_info.feed_target_names,
                            test_program_info.fetch_targets)

                        _logger.info(
                            "epoch: {}, batch: {} metric of compressed model is: {}".
                            format(epoch_id, batch_id, metric))
                        if metric > best_metric:
                            paddle.static.save(
                                program=test_program_info.program._program,
                                model_path=os.path.join(self.save_dir,
                                                        'best_model'))
                        if self.train_config.target_metric is not None:
                            if metric > float(self.train_config.target_metric):
                                return

                    else:
                        raise NotImplementedError(
                            "Please support eval function")

        if 'qat' in self._strategy:
            ### TODO: load best model to save
            float_program, int8_program = convert(test_program_info.program._program, self._places, self._quant_config, \
                                          scope=paddle.static.global_scope(), \
                                          save_int8=True)
            test_program_info.program = float_program
        return test_program_info

    def _save_model(self, test_program_info):
        test_program = test_program_info.program._program if isinstance(
            test_program_info.program,
            paddle.static.CompiledProgram) else test_program_info.program
        feed_vars = []
        for name in test_program_info.feed_target_names:
            for var in test_program.list_vars():
                if var.name == name:
                    feed_vars.append(var)
                    break
        assert len(feed_vars) > 0, "can not find feed vars in quant program"
        paddle.static.save_inference_model(
            path_prefix=os.path.join(self.save_dir, 'final_model'),
            feed_vars=feed_vars,
            fetch_vars=test_program_info.fetch_targets,
            executor=self._exe,
            program=test_program)