|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import inspect |
|
from functools import partial |
|
|
|
|
|
class Registry(object): |
|
def __init__(self, name): |
|
self._name = name |
|
self._module_dict = dict() |
|
|
|
def __repr__(self): |
|
format_str = self.__class__.__name__ + "(name={}, items={})".format( |
|
self._name, list(self._module_dict.keys()) |
|
) |
|
return format_str |
|
|
|
def __len__(self): |
|
return len(self._module_dict) |
|
|
|
@property |
|
def name(self): |
|
return self._name |
|
|
|
@property |
|
def module_dict(self): |
|
return self._module_dict |
|
|
|
def get(self, key): |
|
return self._module_dict.get(key, None) |
|
|
|
def registe_with_name(self, module_name=None, force=False): |
|
return partial(self.register, module_name=module_name, force=force) |
|
|
|
def register(self, module_build_function, module_name=None, force=False): |
|
"""Register a module build function. |
|
Args: |
|
module (:obj:`nn.Module`): Module to be registered. |
|
""" |
|
if not inspect.isfunction(module_build_function): |
|
raise TypeError( |
|
"module_build_function must be a function, but got {}".format( |
|
type(module_build_function) |
|
) |
|
) |
|
if module_name is None: |
|
module_name = module_build_function.__name__ |
|
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_build_function |
|
|
|
return module_build_function |
|
|
|
|
|
MODULE_BUILD_FUNCS = Registry("model build functions") |
|
|
