quant.py 2.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455
  1. # Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from __future__ import absolute_import, division, print_function
  15. import paddle
  16. from paddlex.ppcls.utils import logger
  17. QUANT_CONFIG = {
  18. # weight preprocess type, default is None and no preprocessing is performed.
  19. 'weight_preprocess_type': None,
  20. # activation preprocess type, default is None and no preprocessing is performed.
  21. 'activation_preprocess_type': None,
  22. # weight quantize type, default is 'channel_wise_abs_max'
  23. 'weight_quantize_type': 'channel_wise_abs_max',
  24. # activation quantize type, default is 'moving_average_abs_max'
  25. 'activation_quantize_type': 'moving_average_abs_max',
  26. # weight quantize bit num, default is 8
  27. 'weight_bits': 8,
  28. # activation quantize bit num, default is 8
  29. 'activation_bits': 8,
  30. # data type after quantization, such as 'uint8', 'int8', etc. default is 'int8'
  31. 'dtype': 'int8',
  32. # window size for 'range_abs_max' quantization. default is 10000
  33. 'window_size': 10000,
  34. # The decay coefficient of moving average, default is 0.9
  35. 'moving_rate': 0.9,
  36. # for dygraph quantization, layers of type in quantizable_layer_type will be quantized
  37. 'quantizable_layer_type': ['Conv2D', 'Linear'],
  38. }
  39. def get_quaner(config, model):
  40. if config.get("Slim", False) and config["Slim"].get("quant", False):
  41. from paddleslim.dygraph.quant import QAT
  42. assert config["Slim"]["quant"]["name"].lower(
  43. ) == 'pact', 'Only PACT quantization method is supported now'
  44. QUANT_CONFIG["activation_preprocess_type"] = "PACT"
  45. quanter = QAT(config=QUANT_CONFIG)
  46. quanter.quantize(model)
  47. logger.info("QAT model summary:")
  48. paddle.summary(model, (1, 3, 224, 224))
  49. else:
  50. quanter = None
  51. return quanter