Yolox export & blade optimize support

configs/detection/yolox/yolox_s_8xb16_300e_voc_jit.py

@@ -0,0 +1,90 @@
+_base_ = './yolox_s_8xb16_300e_coco.py'
+
+# s m l x
+img_scale = (640, 640)
+random_size = (14, 26)
+scale_ratio = (0.1, 2)
+
+# tiny nano without mixup
+# img_scale = (416, 416)
+# random_size = (10, 20)
+# scale_ratio = (0.5, 1.5)
+
+CLASSES = [
+    'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat',
+    'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person',
+    'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor'
+]
+
+# dataset settings
+data_root = 'data/voc/'
+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='MMRandomAffine',
+        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,
+        ratio_range=(0.8, 1.6),
+        pad_val=114.0),
+    dict(
+        type='MMPhotoMetricDistortion',
+        brightness_delta=32,
+        contrast_range=(0.5, 1.5),
+        saturation_range=(0.5, 1.5),
+        hue_delta=18),
+    dict(type='MMRandomFlip', flip_ratio=0.5),
+    dict(type='MMResize', keep_ratio=True),
+    dict(type='MMPad', pad_to_square=True, pad_val=(114.0, 114.0, 114.0)),
+    dict(type='MMNormalize', **img_norm_cfg),
+    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='MMPad', pad_to_square=True, pad_val=(114.0, 114.0, 114.0)),
+    dict(type='MMNormalize', **img_norm_cfg),
+    dict(type='DefaultFormatBundle'),
+    dict(type='Collect', keys=['img'])
+]
+
+train_dataset = dict(
+    type='DetImagesMixDataset',
+    data_source=dict(
+        type='DetSourceVOC',
+        path=data_root + 'ImageSets/Main/train.txt',
+        classes=CLASSES,
+        cache_at_init=True),
+    pipeline=train_pipeline,
+    dynamic_scale=img_scale)
+
+val_dataset = dict(
+    type='DetImagesMixDataset',
+    imgs_per_gpu=2,
+    data_source=dict(
+        type='DetSourceVOC',
+        path=data_root + 'ImageSets/Main/val.txt',
+        classes=CLASSES,
+        cache_at_init=True),
+    pipeline=test_pipeline,
+    dynamic_scale=None,
+    label_padding=False)
+
+data = dict(
+    imgs_per_gpu=16, workers_per_gpu=4, train=train_dataset, val=val_dataset)
+
+# # evaluation
+eval_pipelines = [
+    dict(
+        mode='test',
+        data=data['val'],
+        evaluators=[dict(type='CocoDetectionEvaluator', classes=CLASSES)],
+    )
+]
+
+export = dict(use_jit=True, export_blade=False)

easycv/apis/export.py

@@ -9,6 +9,7 @@
 from easycv.file import io
 from easycv.models import (DINO, MOCO, SWAV, YOLOX, Classification, MoBY,
                            build_model)
+from easycv.toolkit.blade import blade_env_assert, blade_optimize
 from easycv.utils.checkpoint import load_checkpoint
 
 __all__ = ['export']
@@ -151,25 +152,70 @@ def _export_yolox(model, cfg, filename):
         cfg: Config object
         filename (str): filename to save exported models
     """
-    if hasattr(cfg, 'test_pipeline'):
-        # with last pipeline Collect
-        test_pipeline = cfg.test_pipeline
-        print(test_pipeline)
-    else:
-        print('test_pipeline not found, using default preprocessing!')
-        raise ValueError('export model config without test_pipeline')
+    cfg.model.type = 'YOLOXExport'
+    if hasattr(cfg, 'export') and getattr(cfg.export, 'use_jit', False):
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        model_export = build_model(cfg.model)
+        model_export.eval()
+        model_export.to(device)
+        model_export.load_state_dict(model.state_dict())
+        batch_size = cfg.export.get('batch_size', 1)
+        img_scale = cfg.get('img_scale', (640, 640))
+        assert (
+            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)
+
+        # well trained model will generate reasonable result, otherwise, we should change model.test_conf=0.0 to avoid tensor in inference to be empty
+        try:
+            yolox_trace = torch.jit.trace(model_export, input.to(device))
+        except:
+            model_export.test_conf = 0.0
+            yolox_trace = torch.jit.trace(model_export, input.to(device))
+
+        if getattr(cfg.export, 'export_blade', False):
+            blade_config = cfg.export.get('blade_config',
+                                          dict(enable_fp16=True))
+            if blade_env_assert() == True:
+                yolox_blade = blade_optimize(
+                    script_model=model_export,
+                    model=yolox_trace,
+                    inputs=(input.to(device), ),
+                    blade_config=blade_config)
+                with io.open(filename + '.blade', 'wb') as ofile:
+                    torch.jit.save(yolox_blade, ofile)
+                with io.open(filename + '.blade.classnames.json',
+                             'w') as ofile:
+                    json.dump(cfg.CLASSES, ofile)
+            else:
+                logging.warning('Export YoloX predictor with blade failed!')
+
+        with io.open(filename + '.jit', 'wb') as ofile:
+            torch.jit.save(yolox_trace, ofile)
+
+        with io.open(filename + '.jit.classnames.json', 'w') as ofile:
+            json.dump(cfg.CLASSES, ofile)
 
-    config = dict(
-        model=cfg.model,
-        test_pipeline=test_pipeline,
-        CLASSES=cfg.CLASSES,
-    )
+    else:
+        if hasattr(cfg, 'test_pipeline'):
+            # with last pipeline Collect
+            test_pipeline = cfg.test_pipeline
+            print(test_pipeline)
+        else:
+            print('test_pipeline not found, using default preprocessing!')
+            raise ValueError('export model config without test_pipeline')
+
+        config = dict(
+            model=cfg.model,
+            test_pipeline=test_pipeline,
+            CLASSES=cfg.CLASSES,
+        )
 
-    meta = dict(config=json.dumps(config))
-    checkpoint = dict(
-        state_dict=model.state_dict(), meta=meta, author='EasyCV')
-    with io.open(filename, 'wb') as ofile:
-        torch.save(checkpoint, ofile)
+        meta = dict(config=json.dumps(config))
+        checkpoint = dict(
+            state_dict=model.state_dict(), meta=meta, author='EasyCV')
+        with io.open(filename, 'wb') as ofile:
+            torch.save(checkpoint, ofile)
 
 
 def _export_swav(model, cfg, filename):

easycv/models/detection/utils/boxes.py

@@ -43,10 +43,10 @@ def postprocess(prediction, num_classes, conf_thre=0.7, nms_thre=0.45):
 
     output = [None for _ in range(len(prediction))]
     for i, image_pred in enumerate(prediction):
-
         # If none are remaining => process next image
         if not image_pred.size(0):
             continue
+
         # Get score and class with highest confidence
         class_conf, class_pred = torch.max(
             image_pred[:, 5:5 + num_classes], 1, keepdim=True)
@@ -57,9 +57,9 @@ def postprocess(prediction, num_classes, conf_thre=0.7, nms_thre=0.45):
         detections = torch.cat(
             (image_pred[:, :5], class_conf, class_pred.float()), 1)
         detections = detections[conf_mask]
+
         if not detections.size(0):
             continue
-
         if LooseVersion(torchvision.__version__) >= LooseVersion('0.8.0'):
             nms_out_index = torchvision.ops.batched_nms(
                 detections[:, :4], detections[:, 4] * detections[:, 5],

easycv/models/detection/yolox/yolox.py

@@ -62,6 +62,7 @@ def __init__(self,
         self.test_conf = test_conf
         self.nms_thre = nms_thre
         self.test_size = test_size
+        self.traceable = False
 
     def forward_train(self,
                       img: Tensor,
@@ -165,3 +166,28 @@ def forward_compression(self, x):
         outputs = self.head(fpn_outs)
 
         return outputs
+
+
+@MODELS.register_module
+class YOLOXExport(YOLOX):
+
+    def __init__(self, *args, **kwargs):
+        super(YOLOXExport, self).__init__(*args, **kwargs)
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.traceable = True
+        self.backbone = self.backbone.to(device)
+        self.head = self.head.to(device)
+        for param in self.backbone.parameters():
+            param.requires_grad = False
+        for param in self.head.parameters():
+            param.requires_grad = False
+
+    def forward(self, img):
+        with torch.no_grad():
+
+            fpn_outs = self.backbone(img)
+            outputs = self.head(fpn_outs)
+
+            outputs = postprocess(outputs, self.num_classes, self.test_conf,
+                                  self.nms_thre)
+        return outputs