192 lines
7.1 KiB
Python
192 lines
7.1 KiB
Python
# Copyright (c) 2022 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.
|
|
|
|
import os
|
|
import sys
|
|
import numpy as np
|
|
import argparse
|
|
import paddle
|
|
from ppdet.core.workspace import load_config, merge_config
|
|
from ppdet.core.workspace import create
|
|
from ppdet.metrics import COCOMetric, VOCMetric, KeyPointTopDownCOCOEval
|
|
from paddleslim.auto_compression.config_helpers import load_config as load_slim_config
|
|
from paddleslim.auto_compression import AutoCompression
|
|
from post_process import PPYOLOEPostProcess
|
|
from paddleslim.common.dataloader import get_feed_vars
|
|
|
|
|
|
def argsparser():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument(
|
|
'--config_path',
|
|
type=str,
|
|
default=None,
|
|
help="path of compression strategy config.",
|
|
required=True)
|
|
parser.add_argument(
|
|
'--save_dir',
|
|
type=str,
|
|
default='output',
|
|
help="directory to save compressed model.")
|
|
parser.add_argument(
|
|
'--devices',
|
|
type=str,
|
|
default='gpu',
|
|
help="which device used to compress.")
|
|
|
|
return parser
|
|
|
|
|
|
def reader_wrapper(reader, input_list):
|
|
def gen():
|
|
for data in reader:
|
|
in_dict = {}
|
|
if isinstance(input_list, list):
|
|
for input_name in input_list:
|
|
in_dict[input_name] = data[input_name]
|
|
elif isinstance(input_list, dict):
|
|
for input_name in input_list.keys():
|
|
in_dict[input_list[input_name]] = data[input_name]
|
|
yield in_dict
|
|
|
|
return gen
|
|
|
|
|
|
def convert_numpy_data(data, metric):
|
|
data_all = {}
|
|
data_all = {k: np.array(v) for k, v in data.items()}
|
|
if isinstance(metric, VOCMetric):
|
|
for k, v in data_all.items():
|
|
if not isinstance(v[0], np.ndarray):
|
|
tmp_list = []
|
|
for t in v:
|
|
tmp_list.append(np.array(t))
|
|
data_all[k] = np.array(tmp_list)
|
|
else:
|
|
data_all = {k: np.array(v) for k, v in data.items()}
|
|
return data_all
|
|
|
|
|
|
def eval_function(exe, compiled_test_program, test_feed_names, test_fetch_list):
|
|
metric = global_config['metric']
|
|
for batch_id, data in enumerate(val_loader):
|
|
data_all = convert_numpy_data(data, metric)
|
|
data_input = {}
|
|
for k, v in data.items():
|
|
if isinstance(global_config['input_list'], list):
|
|
if k in test_feed_names:
|
|
data_input[k] = np.array(v)
|
|
elif isinstance(global_config['input_list'], dict):
|
|
if k in global_config['input_list'].keys():
|
|
data_input[global_config['input_list'][k]] = np.array(v)
|
|
outs = exe.run(compiled_test_program,
|
|
feed=data_input,
|
|
fetch_list=test_fetch_list,
|
|
return_numpy=False)
|
|
res = {}
|
|
if 'include_nms' in global_config and not global_config['include_nms']:
|
|
if 'arch' in global_config and global_config['arch'] == 'PPYOLOE':
|
|
postprocess = PPYOLOEPostProcess(
|
|
score_threshold=0.01, nms_threshold=0.6)
|
|
else:
|
|
assert "Not support arch={} now.".format(global_config['arch'])
|
|
res = postprocess(np.array(outs[0]), data_all['scale_factor'])
|
|
else:
|
|
for out in outs:
|
|
v = np.array(out)
|
|
if len(v.shape) > 1:
|
|
res['bbox'] = v
|
|
else:
|
|
res['bbox_num'] = v
|
|
|
|
metric.update(data_all, res)
|
|
if batch_id % 100 == 0:
|
|
print('Eval iter:', batch_id)
|
|
metric.accumulate()
|
|
metric.log()
|
|
map_res = metric.get_results()
|
|
metric.reset()
|
|
map_key = 'keypoint' if 'arch' in global_config and global_config[
|
|
'arch'] == 'keypoint' else 'bbox'
|
|
return map_res[map_key][0]
|
|
|
|
|
|
def main():
|
|
global global_config
|
|
all_config = load_slim_config(FLAGS.config_path)
|
|
assert "Global" in all_config, "Key 'Global' not found in config file."
|
|
global_config = all_config["Global"]
|
|
reader_cfg = load_config(global_config['reader_config'])
|
|
|
|
train_loader = create('EvalReader')(reader_cfg['TrainDataset'],
|
|
reader_cfg['worker_num'],
|
|
return_list=True)
|
|
if global_config.get('input_list') is None:
|
|
global_config['input_list'] = get_feed_vars(
|
|
global_config['model_dir'], global_config['model_filename'],
|
|
global_config['params_filename'])
|
|
train_loader = reader_wrapper(train_loader, global_config['input_list'])
|
|
|
|
if 'Evaluation' in global_config.keys() and global_config[
|
|
'Evaluation'] and paddle.distributed.get_rank() == 0:
|
|
eval_func = eval_function
|
|
dataset = reader_cfg['EvalDataset']
|
|
global val_loader
|
|
_eval_batch_sampler = paddle.io.BatchSampler(
|
|
dataset, batch_size=reader_cfg['EvalReader']['batch_size'])
|
|
val_loader = create('EvalReader')(dataset,
|
|
reader_cfg['worker_num'],
|
|
batch_sampler=_eval_batch_sampler,
|
|
return_list=True)
|
|
metric = None
|
|
if reader_cfg['metric'] == 'COCO':
|
|
clsid2catid = {v: k for k, v in dataset.catid2clsid.items()}
|
|
anno_file = dataset.get_anno()
|
|
metric = COCOMetric(
|
|
anno_file=anno_file, clsid2catid=clsid2catid, IouType='bbox')
|
|
elif reader_cfg['metric'] == 'VOC':
|
|
metric = VOCMetric(
|
|
label_list=dataset.get_label_list(),
|
|
class_num=reader_cfg['num_classes'],
|
|
map_type=reader_cfg['map_type'])
|
|
elif reader_cfg['metric'] == 'KeyPointTopDownCOCOEval':
|
|
anno_file = dataset.get_anno()
|
|
metric = KeyPointTopDownCOCOEval(anno_file,
|
|
len(dataset), 17, 'output_eval')
|
|
else:
|
|
raise ValueError("metric currently only supports COCO and VOC.")
|
|
global_config['metric'] = metric
|
|
else:
|
|
eval_func = None
|
|
|
|
ac = AutoCompression(
|
|
model_dir=global_config["model_dir"],
|
|
model_filename=global_config["model_filename"],
|
|
params_filename=global_config["params_filename"],
|
|
save_dir=FLAGS.save_dir,
|
|
config=all_config,
|
|
train_dataloader=train_loader,
|
|
eval_callback=eval_func)
|
|
ac.compress()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
paddle.enable_static()
|
|
parser = argsparser()
|
|
FLAGS = parser.parse_args()
|
|
assert FLAGS.devices in ['cpu', 'gpu', 'xpu', 'npu']
|
|
paddle.set_device(FLAGS.devices)
|
|
|
|
main()
|