From f4da38a7336e77b39bda793ef4675cdd3e133947 Mon Sep 17 00:00:00 2001
From: Kamran Melikov <melikovk@gmail.com>
Date: Mon, 20 Jan 2020 02:16:24 -0500
Subject: [PATCH] Add ability to overwite existing module in Registry (#1982)
Changes to be committed:
modified: mmdet/utils/registry.py
---
mmdet/utils/registry.py | 11 +++++++----
1 file changed, 7 insertions(+), 4 deletions(-)
diff --git a/mmdet/utils/registry.py b/mmdet/utils/registry.py
index a1cc87d..4ad9f87 100644
--- a/mmdet/utils/registry.py
+++ b/mmdet/utils/registry.py
@@ -1,4 +1,5 @@
import inspect
+from functools import partial
import mmcv
@@ -25,7 +26,7 @@ class Registry(object):
def get(self, key):
return self._module_dict.get(key, None)
- def _register_module(self, module_class):
+ def _register_module(self, module_class, force=False):
"""Register a module.
Args:
@@ -35,13 +36,15 @@ class Registry(object):
raise TypeError('module must be a class, but got {}'.format(
type(module_class)))
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(
module_name, self.name))
self._module_dict[module_name] = module_class
- def register_module(self, cls):
- self._register_module(cls)
+ def register_module(self, cls=None, force=False):
+ if cls is None:
+ return partial(self.register_module, force=force)
+ self._register_module(cls, force=force)
return cls
--
GitLab