27 lines
857 B
Python
27 lines
857 B
Python
import importlib
|
|
|
|
|
|
def count_params(model, verbose=False):
|
|
total_params = sum(p.numel() for p in model.parameters())
|
|
if verbose:
|
|
print(f"{model.__class__.__name__} has {total_params*1.e-6:.2f} M params.")
|
|
return total_params
|
|
|
|
|
|
def instantiate_from_config(config):
|
|
if not "target" in config:
|
|
if config == '__is_first_stage__':
|
|
return None
|
|
elif config == "__is_unconditional__":
|
|
return None
|
|
raise KeyError("Expected key `target` to instantiate.")
|
|
return get_obj_from_str(config["target"])(**config.get("params", dict()))
|
|
|
|
|
|
def get_obj_from_str(string, reload=False):
|
|
module, cls = string.rsplit(".", 1)
|
|
if reload:
|
|
module_imp = importlib.import_module(module)
|
|
importlib.reload(module_imp)
|
|
return getattr(importlib.import_module(module, package=None), cls)
|