diff --git a/configs/classification/imagenet/resnet/imagenet_resnet50_jpg.py b/configs/classification/imagenet/resnet/imagenet_resnet50_jpg.py index 99f2d169..b678d31b 100644 --- a/configs/classification/imagenet/resnet/imagenet_resnet50_jpg.py +++ b/configs/classification/imagenet/resnet/imagenet_resnet50_jpg.py @@ -6,3 +6,6 @@ depth=50, out_indices=[4], # 0: conv-1, x: stage-x norm_cfg=dict(type='BN'))) + +checkpoint_sync_export = True +export = dict(export_type='raw', export_neck=True) diff --git a/easycv/apis/export.py b/easycv/apis/export.py index 8abc2e76..85424a62 100644 --- a/easycv/apis/export.py +++ b/easycv/apis/export.py @@ -157,6 +157,37 @@ def _get_blade_model(): torch.jit.save(blade_model, ofile) +def _export_onnx_cls(model, model_config, cfg, filename, meta): + + if model_config['backbone'].get( + 'type', None) == 'ResNet' and model_config['backbone'].get( + 'depth', None) == 50: + # save json config for test_pipline and class + with io.open( + filename + + '.config.json' if filename.endswith('onnx') else filename + + '.onnx.config.json', 'w') as ofile: + json.dump(meta, ofile) + + device = 'cuda' if torch.cuda.is_available() else 'cpu' + model.eval() + model.to(device) + img_size = int(cfg.image_size2) + x_input = torch.randn((1, 3, img_size, img_size)).to(device) + torch.onnx.export( + model, + (x_input, 'onnx'), + filename if filename.endswith('onnx') else filename + '.onnx', + export_params=True, + opset_version=12, + do_constant_folding=True, + input_names=['input'], + output_names=['output'], + ) + else: + raise ValueError('Only support export onnx model for ResNet now!') + + def _export_cls(model, cfg, filename): """ export cls (cls & metric learning)model and preprocess config @@ -170,6 +201,7 @@ def _export_cls(model, cfg, filename): else: export_cfg = dict(export_neck=False) + export_type = export_cfg.get('export_type', 'raw') export_neck = export_cfg.get('export_neck', True) label_map_path = cfg.get('label_map_path', None) class_list = None @@ -232,9 +264,14 @@ def _export_cls(model, cfg, filename): if export_neck and (k.startswith('neck') or k.startswith('head')): state_dict[k] = v - checkpoint = dict(state_dict=state_dict, meta=meta, author='EasyCV') - with io.open(filename, 'wb') as ofile: - torch.save(checkpoint, ofile) + if export_type == 'raw': + checkpoint = dict(state_dict=state_dict, meta=meta, author='EasyCV') + with io.open(filename, 'wb') as ofile: + torch.save(checkpoint, ofile) + elif export_type == 'onnx': + _export_onnx_cls(model, model_config, cfg, filename, config) + else: + raise ValueError('Only support export onnx/raw model!') def _export_yolox(model, cfg, filename): diff --git a/easycv/models/classification/classification.py b/easycv/models/classification/classification.py index 57c32c08..479b4d1e 100644 --- a/easycv/models/classification/classification.py +++ b/easycv/models/classification/classification.py @@ -151,6 +151,20 @@ def forward_backbone(self, img: torch.Tensor) -> List[torch.Tensor]: x = self.backbone(img) return x + def forward_onnx(self, img: torch.Tensor) -> Dict[str, torch.Tensor]: + """ + forward_onnx means generate prob from image only support one neck + one head + """ + x = self.forward_backbone(img) # tuple + + # if self.neck_num > 0: + if hasattr(self, 'neck_0'): + x = self.neck_0([i for i in x]) + + out = self.head_0(x)[0].cpu() + out = self.activate_fn(out) + return out + @torch.jit.unused def forward_train(self, img, gt_labels) -> Dict[str, torch.Tensor]: """ @@ -290,6 +304,9 @@ def forward( return self.forward_test_label(img, gt_labels) else: return self.forward_test(img) + elif mode == 'onnx': + return self.forward_onnx(img) + elif mode == 'extract': rd = self.forward_feature(img) rv = {} diff --git a/easycv/predictors/classifier.py b/easycv/predictors/classifier.py index facd9feb..890dd9fd 100644 --- a/easycv/predictors/classifier.py +++ b/easycv/predictors/classifier.py @@ -1,5 +1,7 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +import glob import math +import os import numpy as np import torch @@ -7,11 +9,18 @@ from easycv.file import io from easycv.framework.errors import ValueError +from easycv.utils.checkpoint import load_checkpoint from easycv.utils.misc import deprecated from .base import InputProcessor, OutputProcessor, Predictor, PredictorV2 from .builder import PREDICTORS +# onnx specific +def onnx_to_numpy(tensor): + return tensor.detach().cpu().numpy( + ) if tensor.requires_grad else tensor.cpu().numpy() + + class ClsInputProcessor(InputProcessor): """Process inputs for classification models. @@ -146,6 +155,20 @@ def __init__(self, self.pil_input = pil_input self.label_map_path = label_map_path + if model_path.endswith('onnx'): + self.model_type = 'onnx' + pwd_model = os.path.dirname(model_path) + raw_model = glob.glob( + os.path.join(pwd_model, '*.onnx.config.json')) + if len(raw_model) != 0: + config_file = raw_model[0] + else: + assert len( + raw_model + ) == 0, 'Please have a file with the .onnx.config.json extension in your directory' + else: + self.model_type = 'raw' + if self.pil_input: mode = 'RGB' super(ClassificationPredictor, self).__init__( @@ -186,6 +209,41 @@ def get_output_processor(self): return ClsOutputProcessor(topk=self.topk, label_map=self.label_map) + def prepare_model(self): + """Build model from config file by default. + If the model is not loaded from a configuration file, e.g. torch jit model, you need to reimplement it. + """ + if self.model_type == 'raw': + model = self._build_model() + model.to(self.device) + model.eval() + load_checkpoint(model, self.model_path, map_location='cpu') + return model + else: + import onnxruntime + if onnxruntime.get_device() == 'GPU': + onnx_model = onnxruntime.InferenceSession( + self.model_path, providers=['CUDAExecutionProvider']) + else: + onnx_model = onnxruntime.InferenceSession(self.model_path) + + return onnx_model + + def model_forward(self, inputs): + """Model forward. + If you need refactor model forward, you need to reimplement it. + """ + with torch.no_grad(): + if self.model_type == 'raw': + outputs = self.model(**inputs, mode='test') + else: + outputs = self.model.run(None, { + self.model.get_inputs()[0].name: + onnx_to_numpy(inputs['img']) + })[0] + outputs = dict(prob=torch.from_numpy(outputs)) + return outputs + try: from easy_vision.python.inference.predictor import PredictorInterface diff --git a/easycv/version.py b/easycv/version.py index 5958af20..61c9d5c9 100644 --- a/easycv/version.py +++ b/easycv/version.py @@ -2,5 +2,5 @@ # GENERATED VERSION FILE # TIME: Thu Nov 5 14:17:50 2020 -__version__ = '0.11.6' -short_version = '0.11.6' +__version__ = '0.11.7' +short_version = '0.11.7' diff --git a/tests/test_apis/test_export.py b/tests/test_apis/test_export.py index 7784f3b2..6b1abb80 100644 --- a/tests/test_apis/test_export.py +++ b/tests/test_apis/test_export.py @@ -116,6 +116,7 @@ def test_export_cls_syncbn(self): cfg = mmcv_config_fromfile(config_file) cfg_options = { 'model.backbone.norm_cfg.type': 'SyncBN', + 'export.export_type': 'raw' } if cfg_options is not None: cfg.merge_from_dict(cfg_options) @@ -210,6 +211,27 @@ def test_export_stgcn_jit(self): self.assertTrue(os.path.exists(filename + '.jit')) + def test_export_resnet_onnx(self): + + ckpt_path = PRETRAINED_MODEL_RESNET50 + + easycv_dir = os.path.dirname(easycv.__file__) + + if os.path.exists(os.path.join(easycv_dir, 'configs')): + config_dir = os.path.join(easycv_dir, 'configs') + else: + config_dir = os.path.join(os.path.dirname(easycv_dir), 'configs') + config_file = os.path.join( + config_dir, + 'classification/imagenet/resnet/imagenet_resnet50_jpg.py') + + with tempfile.TemporaryDirectory() as tmpdir: + cfg = mmcv_config_fromfile(config_file) + cfg.export.export_type = 'onnx' + filename = os.path.join(tmpdir, 'imagenet_resnet50') + export(cfg, ckpt_path, filename) + self.assertTrue(os.path.exists(filename + '.onnx')) + if __name__ == '__main__': unittest.main() diff --git a/tests/test_predictors/test_classifier.py b/tests/test_predictors/test_classifier.py index e27844c4..6c380bd9 100644 --- a/tests/test_predictors/test_classifier.py +++ b/tests/test_predictors/test_classifier.py @@ -11,7 +11,8 @@ from easycv.predictors.classifier import ClassificationPredictor from easycv.utils.test_util import clean_up, get_tmp_dir from tests.ut_config import (PRETRAINED_MODEL_RESNET50_WITHOUTHEAD, - IMAGENET_LABEL_TXT, TEST_IMAGES_DIR) + IMAGENET_LABEL_TXT, TEST_IMAGES_DIR, + PRETRAINED_MODEL_RESNET50_ONNX_WITHOUTHEAD) class ClassificationPredictorTest(unittest.TestCase): @@ -33,6 +34,17 @@ def test_single(self): self.assertListEqual(results['class_name'], ['"Persian cat",']) self.assertEqual(len(results['class_probs']), 1000) + def test_onnx_single(self): + checkpoint = PRETRAINED_MODEL_RESNET50_ONNX_WITHOUTHEAD + predict_op = ClassificationPredictor(model_path=checkpoint) + + img_path = os.path.join(TEST_IMAGES_DIR, 'catb.jpg') + + results = predict_op([img_path])[0] + self.assertListEqual(results['class'], [578]) + self.assertListEqual(results['class_name'], ['gown']) + self.assertEqual(len(results['class_probs']), 1000) + def test_batch(self): checkpoint = PRETRAINED_MODEL_RESNET50_WITHOUTHEAD config_file = 'configs/classification/imagenet/resnet/imagenet_resnet50_jpg.py' diff --git a/tests/test_predictors/test_pose_predictor.py b/tests/test_predictors/test_pose_predictor.py index 2838addb..1cf01580 100644 --- a/tests/test_predictors/test_pose_predictor.py +++ b/tests/test_predictors/test_pose_predictor.py @@ -54,10 +54,10 @@ def _base_test(self, predictor): assert_array_almost_equal( result0['bbox'], - np.array([[352.3085, 59.00325, 691.4247, 511.15814, 1.], - [10.511196, 177.74883, 101.824326, 299.49966, 1.], - [224.82036, 114.439865, 312.51306, 231.36348, 1.], - [200.71407, 114.716736, 337.17535, 296.6651, 1.]], + np.array([[438.9, 59., 604.8, 511.2, 0.9], + [10.5, 179.6, 101.8, 297.7, 0.9], + [229.6, 114.4, 307.8, 231.4, 0.6], + [229.4, 114.7, 308.5, 296.7, 0.6]], dtype=np.float32), decimal=1) vis_result = predictor.show_result(img1, result0) @@ -92,10 +92,10 @@ def _base_test(self, predictor): assert_array_almost_equal( result1['bbox'][:4], - np.array([[436.23096, 214.72766, 584.26013, 412.09985, 1.], - [43.990044, 91.04126, 164.28406, 251.43329, 1.], - [127.44148, 100.38604, 254.219, 269.42273, 1.], - [190.08075, 117.31801, 311.22394, 278.8423, 1.]], + np.array([[470.6, 214.7, 549.9, 412.1, 0.9], + [71.6, 91., 136.7, 251.4, 0.9], + [159.7, 100.4, 221.9, 269.4, 0.9], + [219.4, 117.3, 281.9, 278.8, 0.9]], dtype=np.float32), decimal=1) vis_result = predictor.show_result(img2, result1) diff --git a/tests/ut_config.py b/tests/ut_config.py index 0c1896b9..ad8b3383 100644 --- a/tests/ut_config.py +++ b/tests/ut_config.py @@ -179,6 +179,9 @@ PRETRAINED_MODEL_RESNET50_WITHOUTHEAD = os.path.join( BASE_LOCAL_PATH, 'pretrained_models/classification/resnet/resnet50_withhead.pth') +PRETRAINED_MODEL_RESNET50_ONNX_WITHOUTHEAD = os.path.join( + BASE_LOCAL_PATH, + 'pretrained_models/classification/resnet/imagenet_resnet50.onnx') PRETRAINED_MODEL_FACEID = os.path.join(BASE_LOCAL_PATH, 'pretrained_models/faceid') PRETRAINED_MODEL_YOLOXS_EXPORT = os.path.join(