mirror of https://github.com/hpcaitech/ColossalAI
[NFC] polish colossalai/builder/builder.py code style (#1265)
parent
db13f96333
commit
94bfd35184
|
@ -6,7 +6,6 @@ import inspect
|
||||||
from colossalai.registry import *
|
from colossalai.registry import *
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
def build_from_config(module, config: dict):
|
def build_from_config(module, config: dict):
|
||||||
"""Returns an object of :class:`module` constructed from `config`.
|
"""Returns an object of :class:`module` constructed from `config`.
|
||||||
|
|
||||||
|
@ -46,23 +45,20 @@ def build_from_registry(config, registry: Registry):
|
||||||
Raises:
|
Raises:
|
||||||
Exception: Raises an Exception if an error occurred when building from registry.
|
Exception: Raises an Exception if an error occurred when building from registry.
|
||||||
"""
|
"""
|
||||||
config_ = config.copy() # keep the original config untouched
|
config_ = config.copy() # keep the original config untouched
|
||||||
assert isinstance(
|
assert isinstance(registry, Registry), f'Expected type Registry but got {type(registry)}'
|
||||||
registry, Registry), f'Expected type Registry but got {type(registry)}'
|
|
||||||
|
|
||||||
mod_type = config_.pop('type')
|
mod_type = config_.pop('type')
|
||||||
assert registry.has(
|
assert registry.has(mod_type), f'{mod_type} is not found in registry {registry.name}'
|
||||||
mod_type), f'{mod_type} is not found in registry {registry.name}'
|
|
||||||
try:
|
try:
|
||||||
obj = registry.get_module(mod_type)(**config_)
|
obj = registry.get_module(mod_type)(**config_)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(
|
print(f'An error occurred when building {mod_type} from registry {registry.name}', flush=True)
|
||||||
f'An error occurred when building {mod_type} from registry {registry.name}',
|
|
||||||
flush=True)
|
|
||||||
raise e
|
raise e
|
||||||
|
|
||||||
return obj
|
return obj
|
||||||
|
|
||||||
|
|
||||||
def build_gradient_handler(config, model, optimizer):
|
def build_gradient_handler(config, model, optimizer):
|
||||||
"""Returns a gradient handler object of :class:`BaseGradientHandler` constructed from `config`,
|
"""Returns a gradient handler object of :class:`BaseGradientHandler` constructed from `config`,
|
||||||
`model` and `optimizer`.
|
`model` and `optimizer`.
|
||||||
|
|
Loading…
Reference in New Issue