mirrors
/
EasyCV
mirror of https://github.com/alibaba/EasyCV.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Features/self test onnx (#330) 

add yolox onnx export method
pull/331/head^2^2
gulou 2023-10-31 14:56:08 +08:00 committed by GitHub
parent db33ced143
commit 8c3ba59aaf
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 101 additions and 52 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

78
configs/config_templates/yolox_itag.py
View File

@ -49,14 +49,14 @@ img_norm_cfg = dict(
mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
train_pipeline = [
dict(type='MMMosaic', img_scale='${img_scale}', pad_val=114.0),
dict(type='MMMosaic', img_scale=tuple(img_scale), pad_val=114.0),
dict(
type='MMRandomAffine',
scaling_ratio_range='${scale_ratio}',
border=['-${img_scale}[0] // 2', '-${img_scale}[1] // 2']),
scaling_ratio_range=scale_ratio,
border=[img_scale[0] // 2, img_scale[1] // 2]),
dict(
type='MMMixUp', # s m x l; tiny nano will detele
img_scale='${img_scale}',
img_scale=tuple(img_scale),
ratio_range=(0.8, 1.6),
pad_val=114.0),
dict(
@ -70,45 +70,43 @@ train_pipeline = [
dict(type='MMPad', pad_to_square=True, pad_val=(114.0, 114.0, 114.0)),
dict(
type='MMNormalize',
mean='${img_norm_cfg.mean}',
std='${img_norm_cfg.std}',
to_rgb='${img_norm_cfg.to_rgb}'),
mean=img_norm_cfg['mean'],
std=img_norm_cfg['std'],
to_rgb=img_norm_cfg['to_rgb']),
dict(type='DefaultFormatBundle'),
dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels'])
]
test_pipeline = [
dict(type='MMResize', img_scale='${img_scale}', keep_ratio=True),
dict(type='MMResize', img_scale=img_scale, keep_ratio=True),
dict(type='MMPad', pad_to_square=True, pad_val=(114.0, 114.0, 114.0)),
dict(
type='MMNormalize',
mean='${img_norm_cfg.mean}',
std='${img_norm_cfg.std}',
to_rgb='${img_norm_cfg.to_rgb}'),
mean=img_norm_cfg['mean'],
std=img_norm_cfg['std'],
to_rgb=img_norm_cfg['to_rgb']),
dict(type='DefaultFormatBundle'),
dict(type='Collect', keys=['img'])
]
train_path = 'data/coco/train2017.manifest'
val_path = 'data/coco/val2017.manifest'
train_dataset = dict(
type='DetImagesMixDataset',
data_source=dict(type='DetSourcePAI', path=train_path, classes=CLASSES),
pipeline=train_pipeline,
dynamic_scale=tuple(img_scale))
val_dataset = dict(
type='DetImagesMixDataset',
imgs_per_gpu=2,
data_source=dict(type='DetSourcePAI', path=val_path, classes=CLASSES),
pipeline=test_pipeline,
dynamic_scale=None,
label_padding=False)
data = dict(
imgs_per_gpu=16,
workers_per_gpu=4,
train=dict(
type='DetImagesMixDataset',
data_source=dict(
type='DetSourcePAI',
path='data/coco/train2017.manifest',
classes='${CLASSES}'),
pipeline='${train_pipeline}',
dynamic_scale='${img_scale}'),
val=dict(
type='DetImagesMixDataset',
imgs_per_gpu=2,
data_source=dict(
type='DetSourcePAI',
path='data/coco/val2017.manifest',
classes='${CLASSES}'),
pipeline='${test_pipeline}',
dynamic_scale=None,
label_padding=False))
imgs_per_gpu=16, workers_per_gpu=4, train=train_dataset, val=val_dataset)
# additional hooks
interval = 10
@ -120,14 +118,14 @@ custom_hooks = [
priority=48),
dict(
type='SyncRandomSizeHook',
ratio_range='${random_size}',
img_scale='${img_scale}',
interval='${interval}',
ratio_range=random_size,
img_scale=img_scale,
interval=interval,
priority=48),
dict(
type='SyncNormHook',
num_last_epochs=15,
interval='${interval}',
interval=interval,
priority=48)
]
@ -135,23 +133,23 @@ custom_hooks = [
vis_num = 20
score_thr = 0.5
eval_config = dict(
interval='${interval}',
interval=interval,
gpu_collect=False,
visualization_config=dict(
vis_num='${vis_num}',
score_thr='${score_thr}',
vis_num=vis_num,
score_thr=score_thr,
) # show by TensorboardLoggerHookV2
)
eval_pipelines = [
dict(
mode='test',
data='${data.val}',
data=val_dataset,
evaluators=[dict(type='CocoDetectionEvaluator', classes=CLASSES)],
)
]
checkpoint_config = dict(interval='${interval}')
checkpoint_config = dict(interval=interval)
# optimizer
# basic_lr_per_img = 0.01 / 64.0
optimizer = dict(

31
easycv/apis/export.py
View File

@ -247,10 +247,10 @@ def _export_yolox(model, cfg, filename):
if hasattr(cfg, 'export'):
export_type = getattr(cfg.export, 'export_type', 'raw')
default_export_type_list = ['raw', 'jit', 'blade']
default_export_type_list = ['raw', 'jit', 'blade', 'onnx']
if export_type not in default_export_type_list:
logging.warning(
'YOLOX-PAI only supports the export type as [raw,jit,blade], otherwise we use raw as default'
'YOLOX-PAI only supports the export type as [raw,jit,blade,onnx], otherwise we use raw as default'
)
export_type = 'raw'
@ -276,7 +276,7 @@ def _export_yolox(model, cfg, filename):
len(img_scale) == 2
), 'Export YoloX predictor config contains img_scale must be (int, int) tuple!'
input = 255 * torch.rand((batch_size, 3) + img_scale)
input = 255 * torch.rand((batch_size, 3) + tuple(img_scale))
# assert use_trt_efficientnms only happens when static_opt=True
if static_opt is not True:
@ -355,6 +355,31 @@ def _export_yolox(model, cfg, filename):
json.dump(config, ofile)
if export_type == 'onnx':
with io.open(
filename + '.config.json' if filename.endswith('onnx')
else filename + '.onnx.config.json', 'w') as ofile:
config = dict(
model=cfg.model,
export=cfg.export,
test_pipeline=cfg.test_pipeline,
classes=cfg.CLASSES)
json.dump(config, ofile)
torch.onnx.export(
model,
input.to(device),
filename if filename.endswith('onnx') else filename +
'.onnx',
export_params=True,
opset_version=12,
do_constant_folding=True,
input_names=['input'],
output_names=['output'],
)
if export_type == 'jit':
with io.open(filename + '.jit', 'wb') as ofile:
torch.jit.save(yolox_trace, ofile)

37
easycv/predictors/detector.py
View File

@ -23,6 +23,12 @@ except Exception:
from .interface import PredictorInterface
# 将张量转化为ndarray格式
def onnx_to_numpy(tensor):
return tensor.detach().cpu().numpy(
) if tensor.requires_grad else tensor.cpu().numpy()
class DetInputProcessor(InputProcessor):
def build_processor(self):
@ -349,9 +355,11 @@ class YoloXPredictor(DetectionPredictor):
self.model_type = 'jit'
elif model_path.endswith('blade'):
self.model_type = 'blade'
elif model_path.endswith('onnx'):
self.model_type = 'onnx'
else:
self.model_type = 'raw'
assert self.model_type in ['raw', 'jit', 'blade']
assert self.model_type in ['raw', 'jit', 'blade', 'onnx']
if self.model_type == 'blade' or self.use_trt_efficientnms:
import torch_blade
@ -381,8 +389,16 @@ class YoloXPredictor(DetectionPredictor):
def _build_model(self):
if self.model_type != 'raw':
with io.open(self.model_path, 'rb') as infile:
model = torch.jit.load(infile, self.device)
if self.model_type != 'onnx':
with io.open(self.model_path, 'rb') as infile:
model = torch.jit.load(infile, self.device)
else:
import onnxruntime
if onnxruntime.get_device() == 'GPU':
model = onnxruntime.InferenceSession(
self.model_path, providers=['CUDAExecutionProvider'])
else:
model = onnxruntime.InferenceSession(self.model_path)
else:
from easycv.utils.misc import reparameterize_models
model = super()._build_model()
@ -394,8 +410,9 @@ class YoloXPredictor(DetectionPredictor):
If the model is not loaded from a configuration file, e.g. torch jit model, you need to reimplement it.
"""
model = self._build_model()
model.to(self.device)
model.eval()
if self.model_type != 'onnx':
model.to(self.device)
model.eval()
if self.model_type == 'raw':
load_checkpoint(model, self.model_path, map_location='cpu')
return model
@ -406,7 +423,15 @@ class YoloXPredictor(DetectionPredictor):
"""
if self.model_type != 'raw':
with torch.no_grad():
outputs = self.model(inputs['img'])
if self.model_type != 'onnx':
outputs = self.model(inputs['img'])
else:
outputs = self.model.run(
None, {
self.model.get_inputs()[0].name:
onnx_to_numpy(inputs['img'])
})[0]
outputs = torch.from_numpy(outputs)
outputs = {'results': outputs} # convert to dict format
else:
outputs = super().model_forward(inputs)

1
requirements/runtime.txt
View File

@ -13,6 +13,7 @@ lmdb
numba
numpy
nuscenes-devkit
onnxruntime
opencv-python
oss2
packaging

6
tests/test_tools/test_predict.py
View File

@ -83,12 +83,12 @@ class PredictTest(unittest.TestCase):
oss_config = get_oss_config()
ak_id = oss_config['ak_id']
ak_secret = oss_config['ak_secret']
hosts = oss_config['hosts'] + ['oss-cn-hangzhou.aliyuncs.com']
hosts = oss_config['hosts']
hosts = ','.join(_ for _ in hosts)
buckets = oss_config['buckets'] + ['easycv']
buckets = oss_config['buckets']
buckets = ','.join(_ for _ in buckets)
input_file = 'oss://easycv/data/small_test_data/test_images/http_image_list.txt'
input_file = 'oss://pai-vision-data-hz/unittest/local_backup/easycv_nfs/data/test_images/http_image_list.txt'
output_file = tempfile.NamedTemporaryFile('w').name
cmd = f'PYTHONPATH=. python tools/predict.py \
--input_file {input_file} \