export_utils.py 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196
  1. # Copyright (c) 2020 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
  15. from __future__ import division
  16. from __future__ import print_function
  17. import os
  18. import yaml
  19. from collections import OrderedDict
  20. import paddle
  21. from paddlex.ppdet.data.source.category import get_categories
  22. from paddlex.ppdet.utils.logger import setup_logger
  23. logger = setup_logger('ppdet.engine')
  24. # Global dictionary
  25. TRT_MIN_SUBGRAPH = {
  26. 'YOLO': 3,
  27. 'SSD': 60,
  28. 'RCNN': 40,
  29. 'RetinaNet': 40,
  30. 'S2ANet': 80,
  31. 'EfficientDet': 40,
  32. 'Face': 3,
  33. 'TTFNet': 60,
  34. 'FCOS': 16,
  35. 'SOLOv2': 60,
  36. 'HigherHRNet': 3,
  37. 'HRNet': 3,
  38. 'DeepSORT': 3,
  39. 'ByteTrack': 10,
  40. 'JDE': 10,
  41. 'FairMOT': 5,
  42. 'GFL': 16,
  43. 'PicoDet': 3,
  44. 'CenterNet': 5,
  45. 'TOOD': 5,
  46. 'YOLOX': 8,
  47. }
  48. KEYPOINT_ARCH = ['HigherHRNet', 'TopDownHRNet']
  49. MOT_ARCH = ['DeepSORT', 'JDE', 'FairMOT', 'ByteTrack']
  50. def _prune_input_spec(input_spec, program, targets):
  51. # try to prune static program to figure out pruned input spec
  52. # so we perform following operations in static mode
  53. paddle.enable_static()
  54. pruned_input_spec = [{}]
  55. program = program.clone()
  56. program = program._prune(targets=targets)
  57. global_block = program.global_block()
  58. for name, spec in input_spec[0].items():
  59. try:
  60. v = global_block.var(name)
  61. pruned_input_spec[0][name] = spec
  62. except Exception:
  63. pass
  64. paddle.disable_static()
  65. return pruned_input_spec
  66. def _parse_reader(reader_cfg, dataset_cfg, metric, arch, image_shape):
  67. preprocess_list = []
  68. anno_file = dataset_cfg.get_anno()
  69. clsid2catid, catid2name = get_categories(metric, anno_file, arch)
  70. label_list = [str(cat) for cat in catid2name.values()]
  71. fuse_normalize = reader_cfg.get('fuse_normalize', False)
  72. sample_transforms = reader_cfg['sample_transforms']
  73. for st in sample_transforms[1:]:
  74. for key, value in st.items():
  75. p = {'type': key}
  76. if key == 'Resize':
  77. if int(image_shape[1]) != -1:
  78. value['target_size'] = image_shape[1:]
  79. if fuse_normalize and key == 'NormalizeImage':
  80. continue
  81. p.update(value)
  82. preprocess_list.append(p)
  83. batch_transforms = reader_cfg.get('batch_transforms', None)
  84. if batch_transforms:
  85. for bt in batch_transforms:
  86. for key, value in bt.items():
  87. # for deploy/infer, use PadStride(stride) instead PadBatch(pad_to_stride)
  88. if key == 'PadBatch':
  89. preprocess_list.append({
  90. 'type': 'PadStride',
  91. 'stride': value['pad_to_stride']
  92. })
  93. break
  94. return preprocess_list, label_list
  95. def _parse_tracker(tracker_cfg):
  96. tracker_params = {}
  97. for k, v in tracker_cfg.items():
  98. tracker_params.update({k: v})
  99. return tracker_params
  100. def _dump_infer_config(config, path, image_shape, model):
  101. arch_state = False
  102. from paddlex.ppdet.core.config.yaml_helpers import setup_orderdict
  103. setup_orderdict()
  104. use_dynamic_shape = True if image_shape[2] == -1 else False
  105. infer_cfg = OrderedDict({
  106. 'mode': 'paddle',
  107. 'draw_threshold': 0.5,
  108. 'metric': config['metric'],
  109. 'use_dynamic_shape': use_dynamic_shape
  110. })
  111. export_onnx = config.get('export_onnx', False)
  112. infer_arch = config['architecture']
  113. if 'RCNN' in infer_arch and export_onnx:
  114. logger.warning(
  115. "Exporting RCNN model to ONNX only support batch_size = 1")
  116. infer_cfg['export_onnx'] = True
  117. if infer_arch in MOT_ARCH:
  118. if infer_arch == 'DeepSORT':
  119. tracker_cfg = config['DeepSORTTracker']
  120. else:
  121. tracker_cfg = config['JDETracker']
  122. infer_cfg['tracker'] = _parse_tracker(tracker_cfg)
  123. for arch, min_subgraph_size in TRT_MIN_SUBGRAPH.items():
  124. if arch in infer_arch:
  125. infer_cfg['arch'] = arch
  126. infer_cfg['min_subgraph_size'] = min_subgraph_size
  127. arch_state = True
  128. break
  129. if infer_arch == 'YOLOX':
  130. infer_cfg['arch'] = infer_arch
  131. infer_cfg['min_subgraph_size'] = TRT_MIN_SUBGRAPH[infer_arch]
  132. arch_state = True
  133. if not arch_state:
  134. logger.error(
  135. 'Architecture: {} is not supported for exporting model now.\n'.
  136. format(infer_arch) +
  137. 'Please set TRT_MIN_SUBGRAPH in ppdet/engine/export_utils.py')
  138. os._exit(0)
  139. if 'mask_head' in config[config['architecture']] and config[config[
  140. 'architecture']]['mask_head']:
  141. infer_cfg['mask'] = True
  142. label_arch = 'detection_arch'
  143. if infer_arch in KEYPOINT_ARCH:
  144. label_arch = 'keypoint_arch'
  145. if infer_arch in MOT_ARCH:
  146. label_arch = 'mot_arch'
  147. reader_cfg = config['TestMOTReader']
  148. dataset_cfg = config['TestMOTDataset']
  149. else:
  150. reader_cfg = config['TestReader']
  151. dataset_cfg = config['TestDataset']
  152. infer_cfg['Preprocess'], infer_cfg['label_list'] = _parse_reader(
  153. reader_cfg, dataset_cfg, config['metric'], label_arch, image_shape[1:])
  154. if infer_arch == 'PicoDet':
  155. if hasattr(config, 'export') and config['export'].get(
  156. 'post_process',
  157. False) and not config['export'].get('benchmark', False):
  158. infer_cfg['arch'] = 'GFL'
  159. head_name = 'PicoHeadV2' if config['PicoHeadV2'] else 'PicoHead'
  160. infer_cfg['NMS'] = config[head_name]['nms']
  161. # In order to speed up the prediction, the threshold of nms
  162. # is adjusted here, which can be changed in infer_cfg.yml
  163. config[head_name]['nms']["score_threshold"] = 0.3
  164. config[head_name]['nms']["nms_threshold"] = 0.5
  165. infer_cfg['fpn_stride'] = config[head_name]['fpn_stride']
  166. yaml.dump(infer_cfg, open(path, 'w'))
  167. logger.info("Export inference config file to {}".format(
  168. os.path.join(path)))