| | |
| | | import os |
| | | import io |
| | | import shutil |
| | | import logging |
| | | from collections import OrderedDict |
| | | import numpy as np |
| | | |
| | | from omegaconf import DictConfig, OmegaConf |
| | | |
| | | def statistic_model_parameters(model, prefix=None): |
| | | var_dict = model.state_dict() |
| | |
| | | return numel |
| | | |
| | | |
| | | def int2vec(x, vec_dim=8, dtype=np.int): |
| | | def int2vec(x, vec_dim=8, dtype=np.int32): |
| | | b = ('{:0' + str(vec_dim) + 'b}').format(x) |
| | | # little-endian order: lower bit first |
| | | return (np.array(list(b)[::-1]) == '1').astype(dtype) |
| | |
| | | value = value.split(' ') |
| | | ret_dict.append((key, value)) |
| | | return ret_dict |
| | | |
| | | def deep_update(original, update): |
| | | for key, value in update.items(): |
| | | if isinstance(value, dict) and key in original: |
| | | if len(value) == 0: |
| | | original[key] = value |
| | | deep_update(original[key], value) |
| | | else: |
| | | original[key] = value |
| | | |
| | | |
| | | def prepare_model_dir(**kwargs): |
| | | |
| | | |
| | | os.makedirs(kwargs.get("output_dir", "./"), exist_ok=True) |
| | | |
| | | yaml_file = os.path.join(kwargs.get("output_dir", "./"), "config.yaml") |
| | | OmegaConf.save(config=kwargs, f=yaml_file) |
| | | print(kwargs) |
| | | logging.info("config.yaml is saved to: %s", yaml_file) |
| | | |
| | | # model_path = kwargs.get("model_path") |
| | | # if model_path is not None: |
| | | # config_json = os.path.join(model_path, "configuration.json") |
| | | # if os.path.exists(config_json): |
| | | # shutil.copy(config_json, os.path.join(kwargs.get("output_dir", "./"), "configuration.json")) |