Yolox improve with REPConv/ASFF/TOOD

configs/detection/yolox/yolox_l_8xb8_300e_coco.py

@@ -1,4 +1,4 @@
-_base_ = './yolox_s_8xb16_300e_coco.py'
+_base_ = 'configs/detection/yolox/yolox_s_8xb16_300e_coco.py'
 
 # model settings
 model = dict(model_type='l')

easycv/apis/export.py

@@ -13,6 +13,7 @@
 from easycv.file import io
 from easycv.models import (DINO, MOCO, SWAV, YOLOX, Classification, MoBY,
                            build_model)
+from easycv.models.backbones.repvgg_yolox_backbone import RepVGGBlock
 from easycv.utils.bbox_util import scale_coords
 from easycv.utils.checkpoint import load_checkpoint
 
@@ -21,6 +22,24 @@
 ]
 
 
+def reparameterize_models(model):
+    """ reparameterize model for inference, especially for
+            1. rep conv block : merge 3x3 weight 1x1 weights
+        call module switch_to_deploy recursively
+    Args:
+        model: nn.Module
+    """
+    reparameterize_count = 0
+    for layer in model.modules():
+        reparameterize_count += 1
+        if isinstance(layer, RepVGGBlock):
+            layer.switch_to_deploy()
+    logging.info(
+        'export : PAI-export reparameterize_count(RepVGGBlock, ) switch to deploy with {} blocks'
+        .format(reparameterize_count))
+    return model
+
+
 def export(cfg, ckpt_path, filename):
     """ export model for inference
 
@@ -34,6 +53,7 @@ def export(cfg, ckpt_path, filename):
         load_checkpoint(model, ckpt_path, map_location='cpu')
     else:
         cfg.model.backbone.pretrained = False
+    model = reparameterize_models(model)
 
     if isinstance(model, MOCO) or isinstance(model, DINO):
         _export_moco(model, cfg, filename)
@@ -171,6 +191,7 @@ def _export_yolox(model, cfg, filename):
             raise ValueError('`end2end` only support torch1.7.0 and later!')
 
         batch_size = cfg.export.get('batch_size', 1)
+        static_opt = cfg.export.get('static_opt', True)
         img_scale = cfg.get('img_scale', (640, 640))
         assert (
             len(img_scale) == 2
@@ -194,7 +215,6 @@ def _export_yolox(model, cfg, filename):
         # use trace is a litter bit faster than script. But it is not supported in an end2end model.
         if end2end:
             yolox_trace = torch.jit.script(model_export)
-
         else:
             yolox_trace = torch.jit.trace(model_export, input.to(device))
 
@@ -207,13 +227,17 @@ def _export_yolox(model, cfg, filename):
             assert blade_env_assert()
 
             if end2end:
-                input = 255 * torch.rand(img_scale + (3, ))
+                if batch_size == 1:
+                    input = 255 * torch.rand(img_scale + (3, ))
+                else:
+                    input = 255 * torch.rand(img_scale + (3, batch_size))
 
             yolox_blade = blade_optimize(
                 script_model=model,
                 model=yolox_trace,
                 inputs=(input.to(device), ),
-                blade_config=blade_config)
+                blade_config=blade_config,
+                static_opt=static_opt)
 
             with io.open(filename + '.blade', 'wb') as ofile:
                 torch.jit.save(yolox_blade, ofile)
@@ -644,7 +668,15 @@ def __init__(self,
 
         self.example_inputs = example_inputs
         self.preprocess_fn = preprocess_fn
-        self.postprocess_fn = postprocess_fn
+        self.ignore_postprocess = getattr(self.model, 'ignore_postprocess',
+                                          False)
+        if not self.ignore_postprocess:
+            self.postprocess_fn = postprocess_fn
+        else:
+            self.postprocess_fn = None
+        logging.warning(
+            'Model {} ignore_postprocess set to be {} during export !'.format(
+                type(model), self.ignore_postprocess))
         self.trace_model = trace_model
         if self.trace_model:
             self.trace_module()
@@ -669,7 +701,6 @@ def forward(self, image):
                     image = output
 
             model_output = self.model.forward_export(image)
-
             if self.postprocess_fn is not None:
                 model_output = self.postprocess_fn(model_output,
                                                    *preprocess_outputs)

easycv/apis/train.py

@@ -89,6 +89,7 @@ def train_model(model,
 
     # SyncBatchNorm
     open_sync_bn = cfg.get('sync_bn', False)
+
     if open_sync_bn:
         model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
         logger.info('Using SyncBatchNorm()')

easycv/datasets/detection/raw.py

@@ -126,6 +126,9 @@ def visualize(self, results, vis_num=10, score_thr=0.3, **kwargs):
                     dict of image meta info, containing filename, img_shape,
                     origin_img_shape, scale_factor and so on.
         """
+        import copy
+        results = copy.deepcopy(results)
+
         class_names = None
         if hasattr(self.data_source, 'CLASSES'):
             class_names = self.data_source.CLASSES

easycv/hooks/yolox_mode_switch_hook.py

@@ -40,3 +40,6 @@ def before_train_epoch(self, runner):
             train_loader.dataset.update_skip_type_keys(self.skip_type_keys)
             runner.logger.info('Add additional L1 loss now!')
             model.head.use_l1 = True
+
+        if hasattr(runner.model.module, 'epoch_counter'):
+            runner.model.module.epoch_counter = epoch

easycv/models/backbones/__init__.py

@@ -12,6 +12,7 @@
 from .mnasnet import MNASNet
 from .mobilenetv2 import MobileNetV2
 from .pytorch_image_models_wrapper import *
+from .repvgg_yolox_backbone import RepVGGYOLOX
 from .resnest import ResNeSt
 from .resnet import ResNet
 from .resnet_jit import ResNetJIT

easycv/models/backbones/darknet.py

@@ -1,22 +1,22 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 
+import torch
 from torch import nn
 
 from .network_blocks import (BaseConv, CSPLayer, DWConv, Focus, ResLayer,
-                             SPPBottleneck)
+                             SPPBottleneck, SPPFBottleneck)
 
 
 class Darknet(nn.Module):
     # number of blocks from dark2 to dark5.
     depth2blocks = {21: [1, 2, 2, 1], 53: [2, 8, 8, 4]}
 
-    def __init__(
-            self,
-            depth,
-            in_channels=3,
-            stem_out_channels=32,
-            out_features=('dark3', 'dark4', 'dark5'),
-    ):
+    def __init__(self,
+                 depth,
+                 in_channels=3,
+                 stem_out_channels=32,
+                 out_features=('dark3', 'dark4', 'dark5'),
+                 spp_type='spp'):
         """
         Args:
             depth (int): depth of darknet used in model, usually use [21, 53] for this param.
@@ -49,11 +49,18 @@ def __init__(
             *self.make_group_layer(in_channels, num_blocks[2], stride=2))
         in_channels *= 2  # 512
 
-        self.dark5 = nn.Sequential(
-            *self.make_group_layer(in_channels, num_blocks[3], stride=2),
-            *self.make_spp_block([in_channels, in_channels * 2],
-                                 in_channels * 2),
-        )
+        if spp_type == 'spp':
+            self.dark5 = nn.Sequential(
+                *self.make_group_layer(in_channels, num_blocks[3], stride=2),
+                *self.make_spp_block([in_channels, in_channels * 2],
+                                     in_channels * 2),
+            )
+        elif spp_type == 'sppf':
+            self.dark5 = nn.Sequential(
+                *self.make_group_layer(in_channels, num_blocks[3], stride=2),
+                *self.make_sppf_block([in_channels, in_channels * 2],
+                                      in_channels * 2),
+            )
 
     def make_group_layer(self,
                          in_channels: int,
@@ -87,6 +94,23 @@ def make_spp_block(self, filters_list, in_filters):
         ])
         return m
 
+    def make_sppf_block(self, filters_list, in_filters):
+        m = nn.Sequential(*[
+            BaseConv(in_filters, filters_list[0], 1, stride=1, act='lrelu'),
+            BaseConv(
+                filters_list[0], filters_list[1], 3, stride=1, act='lrelu'),
+            SPPBottleneck(
+                in_channels=filters_list[1],
+                out_channels=filters_list[0],
+                activation='lrelu',
+            ),
+            BaseConv(
+                filters_list[0], filters_list[1], 3, stride=1, act='lrelu'),
+            BaseConv(
+                filters_list[1], filters_list[0], 1, stride=1, act='lrelu'),
+        ])
+        return m
+
     def forward(self, x):
         outputs = {}
         x = self.stem(x)
@@ -104,14 +128,13 @@ def forward(self, x):
 
 class CSPDarknet(nn.Module):
 
-    def __init__(
-        self,
-        dep_mul,
-        wid_mul,
-        out_features=('dark3', 'dark4', 'dark5'),
-        depthwise=False,
-        act='silu',
-    ):
+    def __init__(self,
+                 dep_mul,
+                 wid_mul,
+                 out_features=('dark3', 'dark4', 'dark5'),
+                 depthwise=False,
+                 act='silu',
+                 spp_type='spp'):
         super().__init__()
         assert out_features, 'please provide output features of Darknet'
         self.out_features = out_features
@@ -160,19 +183,35 @@ def __init__(
         )
 
         # dark5
-        self.dark5 = nn.Sequential(
-            Conv(base_channels * 8, base_channels * 16, 3, 2, act=act),
-            SPPBottleneck(
-                base_channels * 16, base_channels * 16, activation=act),
-            CSPLayer(
-                base_channels * 16,
-                base_channels * 16,
-                n=base_depth,
-                shortcut=False,
-                depthwise=depthwise,
-                act=act,
-            ),
-        )
+        if spp_type == 'spp':
+            self.dark5 = nn.Sequential(
+                Conv(base_channels * 8, base_channels * 16, 3, 2, act=act),
+                SPPBottleneck(
+                    base_channels * 16, base_channels * 16, activation=act),
+                CSPLayer(
+                    base_channels * 16,
+                    base_channels * 16,
+                    n=base_depth,
+                    shortcut=False,
+                    depthwise=depthwise,
+                    act=act,
+                ),
+            )
+
+        elif spp_type == 'sppf':
+            self.dark5 = nn.Sequential(
+                Conv(base_channels * 8, base_channels * 16, 3, 2, act=act),
+                SPPFBottleneck(
+                    base_channels * 16, base_channels * 16, activation=act),
+                CSPLayer(
+                    base_channels * 16,
+                    base_channels * 16,
+                    n=base_depth,
+                    shortcut=False,
+                    depthwise=depthwise,
+                    act=act,
+                ),
+            )
 
     def forward(self, x):
         outputs = {}