update README.md

README.md

@@ -128,8 +128,9 @@ trainer = CustomTrainer({"generator": generator, "discriminator": discriminator}
 
 ## Distributed training
 
-Easy distributed initializer `homura.init_distributed()` is available. See [imagenet.py](example/imagenet.py) as an example.
+Distributed training is complicated at glance. `homura` has simple APIs, to hide the messy codes for DDP, such as `homura.init_distributed` for the initialization and `homura.is_master` for checking if the process is master or not. 
 
+For details, see `examples/imagenet.py`.
 
 ## Reproducibility
 
@@ -148,6 +149,30 @@ with set_seed(seed):
  other_thing()
 ```
 
+## Registry System
+
+Following major libraries, `homura` also has a simple register system.
+
+```python
+from homura import Registry
+MODEL_REGISTRY = Registry("language_models")
+
+@MODEL_REGISTRY.register
+class Transformer(nn.Module):
+ ...
+
+# or
+
+MODEL_REGISTRY.register(bert_model, 'bert')
+
+# magic
+MODEL_REGISTRY.import_modules(".")
+
+transformer = MODEL_REGISTRY('Transformer')(...)
+# or
+bert = MODEL_REGISTRY('bert', ...)
+```
+
 # Examples
 
 See [examples](examples).
@@ -168,7 +193,7 @@ If you want
 
 * single node multi threads multi gpus
 
-run `python -m torch.distributed.launch --nproc_per_node=$NUM_GPUS imagenet.py distributed.on=true`.
+run `python -m torch.distributed.launch --nproc_per_node=$NUM_GPUS imagenet.py [...]`.
 
 If you want
 

homura/register.py

@@ -43,34 +43,38 @@ def register_from_dict(self,
  def register(self,
  func: Callable = None,
  *,
- name: Optional[str] = None):
+ name: Optional[str] = None
+ ) -> Callable:
  if func is None:
  return functools.partial(self.register, name=name)
 
  _type = self.type
- if _type is not None and not isinstance(_type, types.FunctionType):
+ if _type is not None and not isinstance(func, types.FunctionType):
  if not (isinstance(func, _type) or issubclass(func, _type)):
- raise TypeError(
- f'`func` is expected to be subclass of {_type}.')
+ raise TypeError(f'`func` is expected to be subclass of {_type}.')
 
  if name is None:
  name = func.__name__
 
  if self._registry.get(name) is not None:
- raise KeyError(
- f'Name {name} is already used, try another name!')
+ raise KeyError(f'Name {name} is already used, try another name!')
 
  self._registry[name] = func
  return func
 
  def __call__(self,
- name: str):
+ name: str,
+ *args,
+ **kwargs):
  ret = self._registry.get(name)
  if ret is None:
  _registry = {k.lower(): v for k, v in self._registry.items()}
  ret = _registry.get(name.lower())
  if ret is None:
  raise KeyError(f'Unknown {name} is called!')
+ if len(args) > 0 or len(kwargs) > 0:
+ # if args and kwargs is specified, instantiate
+ ret = ret(*args, **kwargs)
  return ret
 
  @classmethod