Skip to content
Snippets Groups Projects
Commit f4da38a7 authored by Kamran Melikov's avatar Kamran Melikov Committed by Kai Chen
Browse files

Add ability to overwite existing module in Registry (#1982)

Changes to be committed:
modified:   mmdet/utils/registry.py
parent 4b984a7c
No related branches found
No related tags found
No related merge requests found
import inspect import inspect
from functools import partial
import mmcv import mmcv
...@@ -25,7 +26,7 @@ class Registry(object): ...@@ -25,7 +26,7 @@ class Registry(object):
def get(self, key): def get(self, key):
return self._module_dict.get(key, None) return self._module_dict.get(key, None)
def _register_module(self, module_class): def _register_module(self, module_class, force=False):
"""Register a module. """Register a module.
Args: Args:
...@@ -35,13 +36,15 @@ class Registry(object): ...@@ -35,13 +36,15 @@ class Registry(object):
raise TypeError('module must be a class, but got {}'.format( raise TypeError('module must be a class, but got {}'.format(
type(module_class))) type(module_class)))
module_name = module_class.__name__ module_name = module_class.__name__
if module_name in self._module_dict: if not force and module_name in self._module_dict:
raise KeyError('{} is already registered in {}'.format( raise KeyError('{} is already registered in {}'.format(
module_name, self.name)) module_name, self.name))
self._module_dict[module_name] = module_class self._module_dict[module_name] = module_class
def register_module(self, cls): def register_module(self, cls=None, force=False):
self._register_module(cls) if cls is None:
return partial(self.register_module, force=force)
self._register_module(cls, force=force)
return cls return cls
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment