export_utils.py 6.7 KB
Newer Older
K
Kaipeng Deng 已提交
1 2 3 4 5 6 7 8 9 10 11 12
# Copyright (c) 2020 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   
Q
qingqing01 已提交
13 14 15 16 17 18 19 20 21 22
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import yaml
from collections import OrderedDict

G
Guanghua Yu 已提交
23
import paddle
K
Kaipeng Deng 已提交
24
from ppdet.data.source.category import get_categories
Q
qingqing01 已提交
25

K
Kaipeng Deng 已提交
26
from ppdet.utils.logger import setup_logger
27
logger = setup_logger('ppdet.engine')
Q
qingqing01 已提交
28 29 30 31

# Global dictionary
TRT_MIN_SUBGRAPH = {
    'YOLO': 3,
32
    'SSD': 60,
Q
qingqing01 已提交
33 34
    'RCNN': 40,
    'RetinaNet': 40,
35
    'S2ANet': 80,
Q
qingqing01 已提交
36 37
    'EfficientDet': 40,
    'Face': 3,
38
    'TTFNet': 60,
F
Feng Ni 已提交
39
    'FCOS': 16,
Q
qingqing01 已提交
40
    'SOLOv2': 60,
41 42
    'HigherHRNet': 3,
    'HRNet': 3,
43
    'DeepSORT': 3,
44
    'ByteTrack': 10,
45
    'JDE': 10,
46
    'FairMOT': 5,
G
Guanghua Yu 已提交
47 48
    'GFL': 16,
    'PicoDet': 3,
W
wangguanzhong 已提交
49
    'CenterNet': 5,
S
shangliang Xu 已提交
50
    'TOOD': 5,
F
Feng Ni 已提交
51
    'YOLOX': 8,
Q
qingqing01 已提交
52 53
}

54
KEYPOINT_ARCH = ['HigherHRNet', 'TopDownHRNet']
55
MOT_ARCH = ['DeepSORT', 'JDE', 'FairMOT', 'ByteTrack']
56

Q
qingqing01 已提交
57

G
Guanghua Yu 已提交
58 59 60
def _prune_input_spec(input_spec, program, targets):
    # try to prune static program to figure out pruned input spec
    # so we perform following operations in static mode
W
wangguanzhong 已提交
61
    device = paddle.get_device()
G
Guanghua Yu 已提交
62
    paddle.enable_static()
W
wangguanzhong 已提交
63
    paddle.set_device(device)
G
Guanghua Yu 已提交
64 65 66 67 68 69 70 71 72 73
    pruned_input_spec = [{}]
    program = program.clone()
    program = program._prune(targets=targets)
    global_block = program.global_block()
    for name, spec in input_spec[0].items():
        try:
            v = global_block.var(name)
            pruned_input_spec[0][name] = spec
        except Exception:
            pass
W
wangguanzhong 已提交
74
    paddle.disable_static(place=device)
G
Guanghua Yu 已提交
75 76 77
    return pruned_input_spec


K
Kaipeng Deng 已提交
78
def _parse_reader(reader_cfg, dataset_cfg, metric, arch, image_shape):
Q
qingqing01 已提交
79 80 81
    preprocess_list = []

    anno_file = dataset_cfg.get_anno()
K
Kaipeng Deng 已提交
82

Z
zhiboniu 已提交
83
    clsid2catid, catid2name = get_categories(metric, anno_file, arch)
Q
qingqing01 已提交
84 85 86

    label_list = [str(cat) for cat in catid2name.values()]

87
    fuse_normalize = reader_cfg.get('fuse_normalize', False)
Q
qingqing01 已提交
88
    sample_transforms = reader_cfg['sample_transforms']
G
George Ni 已提交
89
    for st in sample_transforms[1:]:
Q
qingqing01 已提交
90 91
        for key, value in st.items():
            p = {'type': key}
G
Guanghua Yu 已提交
92
            if key == 'Resize':
93
                if int(image_shape[1]) != -1:
G
Guanghua Yu 已提交
94
                    value['target_size'] = image_shape[1:]
95
                value['interp'] = value.get('interp', 1)  # cv2.INTER_LINEAR
96 97
            if fuse_normalize and key == 'NormalizeImage':
                continue
Q
qingqing01 已提交
98 99 100 101 102 103
            p.update(value)
            preprocess_list.append(p)
    batch_transforms = reader_cfg.get('batch_transforms', None)
    if batch_transforms:
        for bt in batch_transforms:
            for key, value in bt.items():
104
                # for deploy/infer, use PadStride(stride) instead PadBatch(pad_to_stride)
G
Guanghua Yu 已提交
105
                if key == 'PadBatch':
106 107
                    preprocess_list.append({
                        'type': 'PadStride',
Q
qingqing01 已提交
108 109 110 111
                        'stride': value['pad_to_stride']
                    })
                    break

112
    return preprocess_list, label_list
Q
qingqing01 已提交
113

G
George Ni 已提交
114

115 116 117 118 119
def _parse_tracker(tracker_cfg):
    tracker_params = {}
    for k, v in tracker_cfg.items():
        tracker_params.update({k: v})
    return tracker_params
Q
qingqing01 已提交
120

G
George Ni 已提交
121

K
Kaipeng Deng 已提交
122
def _dump_infer_config(config, path, image_shape, model):
Q
qingqing01 已提交
123 124 125
    arch_state = False
    from ppdet.core.config.yaml_helpers import setup_orderdict
    setup_orderdict()
G
Guanghua Yu 已提交
126
    use_dynamic_shape = True if image_shape[2] == -1 else False
Q
qingqing01 已提交
127
    infer_cfg = OrderedDict({
128
        'mode': 'paddle',
Q
qingqing01 已提交
129 130
        'draw_threshold': 0.5,
        'metric': config['metric'],
131
        'use_dynamic_shape': use_dynamic_shape
Q
qingqing01 已提交
132
    })
133
    export_onnx = config.get('export_onnx', False)
134 135
    export_eb = config.get('export_eb', False)

136

Q
qingqing01 已提交
137
    infer_arch = config['architecture']
138 139 140 141
    if 'RCNN' in infer_arch and export_onnx:
        logger.warning(
            "Exporting RCNN model to ONNX only support batch_size = 1")
        infer_cfg['export_onnx'] = True
142 143
        infer_cfg['export_eb'] = export_eb

Q
qingqing01 已提交
144

145 146 147 148 149 150 151
    if infer_arch in MOT_ARCH:
        if infer_arch == 'DeepSORT':
            tracker_cfg = config['DeepSORTTracker']
        else:
            tracker_cfg = config['JDETracker']
        infer_cfg['tracker'] = _parse_tracker(tracker_cfg)

Q
qingqing01 已提交
152 153 154 155 156 157
    for arch, min_subgraph_size in TRT_MIN_SUBGRAPH.items():
        if arch in infer_arch:
            infer_cfg['arch'] = arch
            infer_cfg['min_subgraph_size'] = min_subgraph_size
            arch_state = True
            break
F
Feng Ni 已提交
158 159 160 161 162 163

    if infer_arch == 'YOLOX':
        infer_cfg['arch'] = infer_arch
        infer_cfg['min_subgraph_size'] = TRT_MIN_SUBGRAPH[infer_arch]
        arch_state = True

Q
qingqing01 已提交
164 165
    if not arch_state:
        logger.error(
G
Guanghua Yu 已提交
166 167 168
            'Architecture: {} is not supported for exporting model now.\n'.
            format(infer_arch) +
            'Please set TRT_MIN_SUBGRAPH in ppdet/engine/export_utils.py')
Q
qingqing01 已提交
169
        os._exit(0)
170 171
    if 'mask_head' in config[config['architecture']] and config[config[
            'architecture']]['mask_head']:
G
Guanghua Yu 已提交
172
        infer_cfg['mask'] = True
173 174 175
    label_arch = 'detection_arch'
    if infer_arch in KEYPOINT_ARCH:
        label_arch = 'keypoint_arch'
176 177 178 179 180 181 182 183 184

    if infer_arch in MOT_ARCH:
        label_arch = 'mot_arch'
        reader_cfg = config['TestMOTReader']
        dataset_cfg = config['TestMOTDataset']
    else:
        reader_cfg = config['TestReader']
        dataset_cfg = config['TestDataset']

185
    infer_cfg['Preprocess'], infer_cfg['label_list'] = _parse_reader(
G
Guanghua Yu 已提交
186
        reader_cfg, dataset_cfg, config['metric'], label_arch, image_shape[1:])
Q
qingqing01 已提交
187

188
    if infer_arch == 'PicoDet':
G
Guanghua Yu 已提交
189 190 191
        if hasattr(config, 'export') and config['export'].get(
                'post_process',
                False) and not config['export'].get('benchmark', False):
192
            infer_cfg['arch'] = 'GFL'
193 194
        head_name = 'PicoHeadV2' if config['PicoHeadV2'] else 'PicoHead'
        infer_cfg['NMS'] = config[head_name]['nms']
195 196
        # In order to speed up the prediction, the threshold of nms 
        # is adjusted here, which can be changed in infer_cfg.yml
197 198 199
        config[head_name]['nms']["score_threshold"] = 0.3
        config[head_name]['nms']["nms_threshold"] = 0.5
        infer_cfg['fpn_stride'] = config[head_name]['fpn_stride']
200

Q
qingqing01 已提交
201 202
    yaml.dump(infer_cfg, open(path, 'w'))
    logger.info("Export inference config file to {}".format(os.path.join(path)))