11 files changed, 470 insertions, 13 deletions

diff --git a/README.md b/README.md
index 5066e034..61b2817f 100644
--- a/README.md
+++ b/README.md
@@ -15,7 +15,9 @@ This repository is released under the following [LICENSE](./LICENSE).
 
 ## Installation
 
-Please refer to [mmdetection install.md](https://github.com/open-mmlab/mmdetection/docs/install.md) for installation and dataset preparation.
+Please refer to [mmdetection install.md](https://github.com/open-mmlab/mmdetection/docs/install.md) for installation and dataset preparation. 
+
+We used the the version **v2.1.0** of mmdetection to test our changes. If you get any issues with the master branch of mmdetection, try checking out that tag.
 
 After installing mmdetection, please install [PyTorch-Jacinto-AI-DevKit](https://bitbucket.itg.ti.com/projects/JACINTO-AI/repos/pytorch-jacinto-ai-devkit/browse/) as our repository uses several components from there - especially to define low complexity models and to Quantization Aware Training (QAT).
 
diff --git a/tools/test.py b/tools/test.py
index e5c36446..803e3860 100644
--- a/tools/test.py
+++ b/tools/test.py
@@ -8,11 +8,11 @@ from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
 from mmcv.runner import get_dist_info, init_dist, load_checkpoint
 from tools.fuse_conv_bn import fuse_module
 
-from mmdet.apis import multi_gpu_test, single_gpu_test
-from mmdet.core import wrap_fp16_model
-from mmdet.datasets import build_dataloader, build_dataset
-from mmdet.models import build_detector
-from mmdet.utils import MMDetQuantTestModule, save_model_proto, mmdet_load_checkpoint
+from xmmdet.apis import multi_gpu_test, single_gpu_test
+from xmmdet.core import wrap_fp16_model
+from xmmdet.datasets import build_dataloader, build_dataset
+from xmmdet.models import build_detector
+from xmmdet.utils import MMDetQuantTestModule, save_model_proto, mmdet_load_checkpoint
 
 from pytorch_jacinto_ai import xnn
 
diff --git a/xmmdet/__init__.py b/xmmdet/__init__.py
index f9d66cc7..325fb75f 100644
--- a/xmmdet/__init__.py
+++ b/xmmdet/__init__.py
@@ -1,2 +1,8 @@
 from mmdet import *
+from .ops import *
+from .core import *
+from .datasets import *
+from .models import *
+from .utils import *
+from .apis import *
 
diff --git a/xmmdet/models/__init__.py b/xmmdet/models/__init__.py
index 608be121..718b897e 100644
--- a/xmmdet/models/__init__.py
+++ b/xmmdet/models/__init__.py
@@ -1 +1,4 @@
 from mmdet.models import *
+from .backbones import *
+from .dense_heads import *
+from .necks import *
diff --git a/xmmdet/models/backbones/resnet.py b/xmmdet/models/backbones/resnet.py
index 08a26836..5048a4e9 100644
--- a/xmmdet/models/backbones/resnet.py
+++ b/xmmdet/models/backbones/resnet.py
@@ -290,7 +290,7 @@ class Bottleneck(nn.Module):
         return out
 
 
-@BACKBONES.register_module()
+@BACKBONES.register_module(force=True)
 class ResNet(nn.Module):
     """ResNet backbone.
 
@@ -619,7 +619,7 @@ class ResNet(nn.Module):
                     m.eval()
 
 
-@BACKBONES.register_module()
+@BACKBONES.register_module(force=True)
 class ResNetV1d(ResNet):
     """ResNetV1d variant described in
     `Bag of Tricks <https://arxiv.org/pdf/1812.01187.pdf>`_.
diff --git a/xmmdet/models/dense_heads/fcos_head.py b/xmmdet/models/dense_heads/fcos_head.py
index 1961245c..27e15c3d 100644
--- a/xmmdet/models/dense_heads/fcos_head.py
+++ b/xmmdet/models/dense_heads/fcos_head.py
@@ -8,7 +8,7 @@ from mmdet.models.builder import HEADS, build_loss
 INF = 1e8
 
 
-@HEADS.register_module()
+@HEADS.register_module(force=True)
 class FCOSHead(nn.Module):
     """Anchor-free head used in `FCOS <https://arxiv.org/abs/1904.01355>`_.
 
diff --git a/xmmdet/models/dense_heads/retina_head.py b/xmmdet/models/dense_heads/retina_head.py
index 68232c45..27031acd 100644
--- a/xmmdet/models/dense_heads/retina_head.py
+++ b/xmmdet/models/dense_heads/retina_head.py
@@ -5,7 +5,7 @@ from mmdet.models.builder import HEADS
 from mmdet.models.dense_heads.anchor_head import AnchorHead
 
 
-@HEADS.register_module()
+@HEADS.register_module(force=True)
 class RetinaHead(AnchorHead):
     """An anchor-based head used in
     `RetinaNet <https://arxiv.org/pdf/1708.02002.pdf>`_.
diff --git a/xmmdet/models/dense_heads/ssd_head.py b/xmmdet/models/dense_heads/ssd_head.py
index d5f4658e..486343f1 100644
--- a/xmmdet/models/dense_heads/ssd_head.py
+++ b/xmmdet/models/dense_heads/ssd_head.py
@@ -11,7 +11,7 @@ from mmdet.models.dense_heads.anchor_head import AnchorHead
 
 
 # TODO: add loss evaluator for SSD
-@HEADS.register_module()
+@HEADS.register_module(force=True)
 class SSDHead(AnchorHead):
 
     def __init__(self,
diff --git a/xmmdet/models/necks/fpn.py b/xmmdet/models/necks/fpn.py
index c55ef05e..f773a01c 100644
--- a/xmmdet/models/necks/fpn.py
+++ b/xmmdet/models/necks/fpn.py
@@ -8,7 +8,7 @@ from mmdet.models.builder import NECKS
 from pytorch_jacinto_ai import xnn
 
 
-@NECKS.register_module()
+@NECKS.register_module(force=True)
 class FPN(nn.Module):
     """
     Feature Pyramid Network.
diff --git a/xmmdet/utils/__init__.py b/xmmdet/utils/__init__.py
index 062baecd..2ce3a842 100644
--- a/xmmdet/utils/__init__.py
+++ b/xmmdet/utils/__init__.py
@@ -1,6 +1,8 @@
 from mmdet.utils import *
+from .flops_counter import get_model_complexity_info
 from .logger import LoggerStream, get_root_logger
-from .runner import MMDetRunner, MMDetNoOptimizerHook
+from .runner import MMDetRunner, MMDetNoOptimizerHook, \
+    mmdet_load_checkpoint, mmdet_save_checkpoint
 from .save_model import save_model_proto
 from .quantize import MMDetQuantTrainModule, MMDetQuantCalibrateModule, \
     MMDetQuantTestModule, is_mmdet_quant_module
diff --git a/xmmdet/utils/flops_counter.py b/xmmdet/utils/flops_counter.py
new file mode 100644
index 00000000..04f27f8b
--- /dev/null
+++ b/xmmdet/utils/flops_counter.py
@@ -0,0 +1,444 @@
+# Modified from flops-counter.pytorch by Vladislav Sovrasov
+# original repo: https://github.com/sovrasov/flops-counter.pytorch
+
+# MIT License
+
+# Copyright (c) 2018 Vladislav Sovrasov
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import sys
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn.modules.batchnorm import _BatchNorm
+from torch.nn.modules.conv import _ConvNd, _ConvTransposeMixin
+from torch.nn.modules.pooling import (_AdaptiveAvgPoolNd, _AdaptiveMaxPoolNd,
+                                      _AvgPoolNd, _MaxPoolNd)
+
+
+def get_model_complexity_info(model,
+                              input_res,
+                              print_per_layer_stat=True,
+                              as_strings=True,
+                              input_constructor=None,
+                              ost=sys.stdout):
+    assert type(input_res) is tuple
+    assert len(input_res) >= 2
+    flops_model = add_flops_counting_methods(model)
+    flops_model.eval().start_flops_count()
+    if input_constructor:
+        input = input_constructor(input_res)
+        _ = flops_model(**input)
+    else:
+        batch = torch.ones(()).new_empty(
+            (1, *input_res),
+            dtype=next(flops_model.parameters()).dtype,
+            device=next(flops_model.parameters()).device)
+        flops_model.forward_dummy(batch)
+
+    if print_per_layer_stat:
+        print_model_with_flops(flops_model, ost=ost)
+    flops_count = flops_model.compute_average_flops_cost()
+    params_count = get_model_parameters_number(flops_model)
+    flops_model.stop_flops_count()
+
+    if as_strings:
+        return flops_to_string(flops_count), params_to_string(params_count)
+
+    return flops_count, params_count
+
+
+def flops_to_string(flops, units='GMac', precision=2):
+    if units is None:
+        if flops // 10**9 > 0:
+            return str(round(flops / 10.**9, precision)) + ' GMac'
+        elif flops // 10**6 > 0:
+            return str(round(flops / 10.**6, precision)) + ' MMac'
+        elif flops // 10**3 > 0:
+            return str(round(flops / 10.**3, precision)) + ' KMac'
+        else:
+            return str(flops) + ' Mac'
+    else:
+        if units == 'GMac':
+            return str(round(flops / 10.**9, precision)) + ' ' + units
+        elif units == 'MMac':
+            return str(round(flops / 10.**6, precision)) + ' ' + units
+        elif units == 'KMac':
+            return str(round(flops / 10.**3, precision)) + ' ' + units
+        else:
+            return str(flops) + ' Mac'
+
+
+def params_to_string(params_num):
+    """converting number to string
+
+    :param float params_num: number
+    :returns str: number
+
+    >>> params_to_string(1e9)
+    '1000.0 M'
+    >>> params_to_string(2e5)
+    '200.0 k'
+    >>> params_to_string(3e-9)
+    '3e-09'
+    """
+    if params_num // 10**6 > 0:
+        return str(round(params_num / 10**6, 2)) + ' M'
+    elif params_num // 10**3:
+        return str(round(params_num / 10**3, 2)) + ' k'
+    else:
+        return str(params_num)
+
+
+def print_model_with_flops(model, units='GMac', precision=3, ost=sys.stdout):
+    total_flops = model.compute_average_flops_cost()
+
+    def accumulate_flops(self):
+        if is_supported_instance(self):
+            return self.__flops__ / (model.__batch_counter__ or 1)
+        else:
+            sum = 0
+            for m in self.children():
+                sum += m.accumulate_flops()
+            return sum
+
+    def flops_repr(self):
+        accumulated_flops_cost = self.accumulate_flops()
+        return ', '.join([
+            flops_to_string(
+                accumulated_flops_cost, units=units, precision=precision),
+            f'{accumulated_flops_cost / total_flops:.3%} MACs',
+            self.original_extra_repr()
+        ])
+
+    def add_extra_repr(m):
+        m.accumulate_flops = accumulate_flops.__get__(m)
+        flops_extra_repr = flops_repr.__get__(m)
+        if m.extra_repr != flops_extra_repr:
+            m.original_extra_repr = m.extra_repr
+            m.extra_repr = flops_extra_repr
+            assert m.extra_repr != m.original_extra_repr
+
+    def del_extra_repr(m):
+        if hasattr(m, 'original_extra_repr'):
+            m.extra_repr = m.original_extra_repr
+            del m.original_extra_repr
+        if hasattr(m, 'accumulate_flops'):
+            del m.accumulate_flops
+
+    model.apply(add_extra_repr)
+    print(model, file=ost)
+    model.apply(del_extra_repr)
+
+
+def get_model_parameters_number(model):
+    params_num = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    return params_num
+
+
+def add_flops_counting_methods(net_main_module):
+    # adding additional methods to the existing module object,
+    # this is done this way so that each function has access to self object
+    net_main_module.start_flops_count = start_flops_count.__get__(
+        net_main_module)
+    net_main_module.stop_flops_count = stop_flops_count.__get__(
+        net_main_module)
+    net_main_module.reset_flops_count = reset_flops_count.__get__(
+        net_main_module)
+    net_main_module.compute_average_flops_cost = \
+        compute_average_flops_cost.__get__(net_main_module)
+
+    net_main_module.reset_flops_count()
+
+    # Adding variables necessary for masked flops computation
+    net_main_module.apply(add_flops_mask_variable_or_reset)
+
+    return net_main_module
+
+
+def compute_average_flops_cost(self):
+    """
+    A method that will be available after add_flops_counting_methods() is
+    called on a desired net object.
+    Returns current mean flops consumption per image.
+    """
+
+    batches_count = (self.__batch_counter__ or 1)
+    flops_sum = 0
+    for module in self.modules():
+        if is_supported_instance(module):
+            flops_sum += module.__flops__
+
+    return flops_sum / batches_count
+
+
+def start_flops_count(self):
+    """
+    A method that will be available after add_flops_counting_methods() is
+    called on a desired net object.
+    Activates the computation of mean flops consumption per image.
+    Call it before you run the network.
+    """
+    add_batch_counter_hook_function(self)
+    self.apply(add_flops_counter_hook_function)
+
+
+def stop_flops_count(self):
+    """
+    A method that will be available after add_flops_counting_methods() is
+    called on a desired net object.
+    Stops computing the mean flops consumption per image.
+    Call whenever you want to pause the computation.
+    """
+    remove_batch_counter_hook_function(self)
+    self.apply(remove_flops_counter_hook_function)
+
+
+def reset_flops_count(self):
+    """
+    A method that will be available after add_flops_counting_methods() is
+    called on a desired net object.
+    Resets statistics computed so far.
+    """
+    add_batch_counter_variables_or_reset(self)
+    self.apply(add_flops_counter_variable_or_reset)
+
+
+def add_flops_mask(module, mask):
+
+    def add_flops_mask_func(module):
+        if isinstance(module, torch.nn.Conv2d):
+            module.__mask__ = mask
+
+    module.apply(add_flops_mask_func)
+
+
+def remove_flops_mask(module):
+    module.apply(add_flops_mask_variable_or_reset)
+
+
+def is_supported_instance(module):
+    for mod in hook_mapping:
+        if issubclass(type(module), mod):
+            return True
+    return False
+
+
+def empty_flops_counter_hook(module, input, output):
+    module.__flops__ += 0
+
+
+def upsample_flops_counter_hook(module, input, output):
+    output_size = output[0]
+    batch_size = output_size.shape[0]
+    output_elements_count = batch_size
+    for val in output_size.shape[1:]:
+        output_elements_count *= val
+    module.__flops__ += int(output_elements_count)
+
+
+def relu_flops_counter_hook(module, input, output):
+    active_elements_count = output.numel()
+    module.__flops__ += int(active_elements_count)
+
+
+def linear_flops_counter_hook(module, input, output):
+    input = input[0]
+    batch_size = input.shape[0]
+    module.__flops__ += int(batch_size * input.shape[1] * output.shape[1])
+
+
+def pool_flops_counter_hook(module, input, output):
+    input = input[0]
+    module.__flops__ += int(np.prod(input.shape))
+
+
+def bn_flops_counter_hook(module, input, output):
+    input = input[0]
+
+    batch_flops = np.prod(input.shape)
+    if module.affine:
+        batch_flops *= 2
+    module.__flops__ += int(batch_flops)
+
+
+def gn_flops_counter_hook(module, input, output):
+    elems = np.prod(input[0].shape)
+    # there is no precise FLOPs estimation of computing mean and variance,
+    # and we just set it 2 * elems: half muladds for computing
+    # means and half for computing vars
+    batch_flops = 3 * elems
+    if module.affine:
+        batch_flops += elems
+    module.__flops__ += int(batch_flops)
+
+
+def deconv_flops_counter_hook(conv_module, input, output):
+    # Can have multiple inputs, getting the first one
+    input = input[0]
+
+    batch_size = input.shape[0]
+    input_height, input_width = input.shape[2:]
+
+    kernel_height, kernel_width = conv_module.kernel_size
+    in_channels = conv_module.in_channels
+    out_channels = conv_module.out_channels
+    groups = conv_module.groups
+
+    filters_per_channel = out_channels // groups
+    conv_per_position_flops = (
+        kernel_height * kernel_width * in_channels * filters_per_channel)
+
+    active_elements_count = batch_size * input_height * input_width
+    overall_conv_flops = conv_per_position_flops * active_elements_count
+    bias_flops = 0
+    if conv_module.bias is not None:
+        output_height, output_width = output.shape[2:]
+        bias_flops = out_channels * batch_size * output_height * output_height
+    overall_flops = overall_conv_flops + bias_flops
+
+    conv_module.__flops__ += int(overall_flops)
+
+
+def conv_flops_counter_hook(conv_module, input, output):
+    # Can have multiple inputs, getting the first one
+    input = input[0]
+
+    batch_size = input.shape[0]
+    output_dims = list(output.shape[2:])
+
+    kernel_dims = list(conv_module.kernel_size)
+    in_channels = conv_module.in_channels
+    out_channels = conv_module.out_channels
+    groups = conv_module.groups
+
+    filters_per_channel = out_channels // groups
+    conv_per_position_flops = np.prod(
+        kernel_dims) * in_channels * filters_per_channel
+
+    active_elements_count = batch_size * np.prod(output_dims)
+
+    if conv_module.__mask__ is not None:
+        # (b, 1, h, w)
+        output_height, output_width = output.shape[2:]
+        flops_mask = conv_module.__mask__.expand(batch_size, 1, output_height,
+                                                 output_width)
+        active_elements_count = flops_mask.sum()
+
+    overall_conv_flops = conv_per_position_flops * active_elements_count
+
+    bias_flops = 0
+
+    if conv_module.bias is not None:
+
+        bias_flops = out_channels * active_elements_count
+
+    overall_flops = overall_conv_flops + bias_flops
+
+    conv_module.__flops__ += int(overall_flops)
+
+
+hook_mapping = {
+    # conv
+    _ConvNd: conv_flops_counter_hook,
+    # deconv
+    _ConvTransposeMixin: deconv_flops_counter_hook,
+    # fc
+    nn.Linear: linear_flops_counter_hook,
+    # pooling
+    _AvgPoolNd: pool_flops_counter_hook,
+    _MaxPoolNd: pool_flops_counter_hook,
+    _AdaptiveAvgPoolNd: pool_flops_counter_hook,
+    _AdaptiveMaxPoolNd: pool_flops_counter_hook,
+    # activation
+    nn.ReLU: relu_flops_counter_hook,
+    nn.PReLU: relu_flops_counter_hook,
+    nn.ELU: relu_flops_counter_hook,
+    nn.LeakyReLU: relu_flops_counter_hook,
+    nn.ReLU6: relu_flops_counter_hook,
+    # normalization
+    _BatchNorm: bn_flops_counter_hook,
+    nn.GroupNorm: gn_flops_counter_hook,
+    # upsample
+    nn.Upsample: upsample_flops_counter_hook,
+}
+
+
+def batch_counter_hook(module, input, output):
+    batch_size = 1
+    if len(input) > 0:
+        # Can have multiple inputs, getting the first one
+        input = input[0]
+        batch_size = len(input)
+    else:
+        print('Warning! No positional inputs found for a module, '
+              'assuming batch size is 1.')
+    module.__batch_counter__ += batch_size
+
+
+def add_batch_counter_variables_or_reset(module):
+    module.__batch_counter__ = 0
+
+
+def add_batch_counter_hook_function(module):
+    if hasattr(module, '__batch_counter_handle__'):
+        return
+
+    handle = module.register_forward_hook(batch_counter_hook)
+    module.__batch_counter_handle__ = handle
+
+
+def remove_batch_counter_hook_function(module):
+    if hasattr(module, '__batch_counter_handle__'):
+        module.__batch_counter_handle__.remove()
+        del module.__batch_counter_handle__
+
+
+def add_flops_counter_variable_or_reset(module):
+    if is_supported_instance(module):
+        module.__flops__ = 0
+
+
+def add_flops_counter_hook_function(module):
+    if is_supported_instance(module):
+        if hasattr(module, '__flops_handle__'):
+            return
+
+        for mod_type, counter_hook in hook_mapping.items():
+            if issubclass(type(module), mod_type):
+                handle = module.register_forward_hook(counter_hook)
+                break
+
+        module.__flops_handle__ = handle
+
+
+def remove_flops_counter_hook_function(module):
+    if is_supported_instance(module):
+        if hasattr(module, '__flops_handle__'):
+            module.__flops_handle__.remove()
+            del module.__flops_handle__
+
+
+# --- Masked flops counting
+# Also being run in the initialization
+def add_flops_mask_variable_or_reset(module):
+    if is_supported_instance(module):
+        module.__mask__ = None