| | 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 |
