| from importlib import import_module |
| from omegaconf import OmegaConf |
| import os |
| from pathlib import Path |
| import shutil |
| from omegaconf import DictConfig |
| from lightning.pytorch.utilities import rank_zero_info |
|
|
| def instantiate(config: DictConfig, instantiate_module=True): |
| """Get arguments from config.""" |
| module = import_module(config.module_name) |
| class_ = getattr(module, config.class_name) |
| if instantiate_module: |
| init_args = {k: v for k, v in config.items() if k not in ["module_name", "class_name"]} |
| return class_(**init_args) |
| else: |
| return class_ |
|
|
| def instantiate_motion_gen(module_name, class_name, cfg, hfstyle=False, **init_args): |
| module = import_module(module_name) |
| class_ = getattr(module, class_name) |
| if hfstyle: |
| config_class = class_.config_class |
| cfg = config_class(config_obj=cfg) |
| return class_(cfg, **init_args) |
| |
| def save_config_and_codes(config, save_dir): |
| os.makedirs(save_dir, exist_ok=True) |
| sanity_check_dir = os.path.join(save_dir, 'sanity_check') |
| os.makedirs(sanity_check_dir, exist_ok=True) |
| with open(os.path.join(sanity_check_dir, f'{config.exp_name}.yaml'), 'w') as f: |
| OmegaConf.save(config, f) |
| current_dir = Path.cwd() |
| for py_file in current_dir.rglob('*.py'): |
| dest_path = Path(sanity_check_dir) / py_file.relative_to(current_dir) |
| dest_path.parent.mkdir(parents=True, exist_ok=True) |
| shutil.copy(py_file, dest_path) |
| |
| def print_model_size(model): |
| total_params = sum(p.numel() for p in model.parameters()) |
| trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) |
| rank_zero_info(f"Total parameters: {total_params:,}") |
| rank_zero_info(f"Trainable parameters: {trainable_params:,}") |
| rank_zero_info(f"Non-trainable parameters: {(total_params - trainable_params):,}") |
| |
| def load_metrics(file_path): |
| metrics = {} |
| with open(file_path, "r") as f: |
| for line in f: |
| key, value = line.strip().split(": ") |
| try: |
| metrics[key] = float(value) |
| except ValueError: |
| metrics[key] = value |
| return metrics |
