Upload directory

models/base/utils.py

@@ -0,0 +1,91 @@
+import itertools
+from typing import List, Optional, Tuple, Union
+import safetensors
+import torch
+from torch import Tensor
+import os
+from pathlib import Path
+from omegaconf import DictConfig, OmegaConf
+
+
+def get_parameter_device(parameter: torch.nn.Module):
+ try:
+ parameters_and_buffers = itertools.chain(parameter.parameters(), parameter.buffers())
+ return next(parameters_and_buffers).device
+ except StopIteration:
+ # For torch.nn.DataParallel compatibility in PyTorch 1.5
+ def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, Tensor]]:
+ tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
+ return tuples
+ gen = parameter._named_members(get_members_fn=find_tensor_attributes)
+ first_tuple = next(gen)
+ return first_tuple[1].device
+
+
+def get_parameter_dtype(parameter: torch.nn.Module):
+ try:
+ params = tuple(parameter.parameters())
+ if len(params) > 0:
+ return params[0].dtype
+
+ buffers = tuple(parameter.buffers())
+ if len(buffers) > 0:
+ return buffers[0].dtype
+
+ except StopIteration:
+ # For torch.nn.DataParallel compatibility in PyTorch 1.5
+
+ def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, Tensor]]:
+ tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
+ return tuples
+
+ gen = parameter._named_members(get_members_fn=find_tensor_attributes)
+ first_tuple = next(gen)
+ return first_tuple[1].dtype
+
+
+def get_parent_directory(save_path: Union[str, os.PathLike]) -> Path:
+ path_obj = Path(save_path)
+ return path_obj.parent
+
+def get_base_name(save_path: Union[str, os.PathLike]) -> str:
+ path_obj = Path(save_path)
+ return path_obj.name
+
+def load_state_dict_from_path(path: Union[str, os.PathLike]):
+ # Load a state dict from a path.
+ if 'safetensors' in path:
+ state_dict = safetensors.torch.load_file(path)
+ else:
+ state_dict = torch.load(path, map_location="cpu")
+ return state_dict
+
+def replace_extension(path, new_extension):
+ if not new_extension.startswith('.'):
+ new_extension = '.' + new_extension
+ return os.path.splitext(path)[0] + new_extension
+
+def make_config_path(save_path):
+ config_path = replace_extension(save_path, '.yaml')
+ return config_path
+
+def save_config(config, config_path):
+ assert isinstance(config, dict) or isinstance(config, DictConfig)
+ os.makedirs(get_parent_directory(config_path), exist_ok=True)
+ if isinstance(config, dict):
+ config = OmegaConf.create(config)
+ OmegaConf.save(config, config_path)
+
+
+def save_state_dict_and_config(state_dict, config, save_path):
+ os.makedirs(get_parent_directory(save_path), exist_ok=True)
+
+ # save config dict
+ config_path = make_config_path(save_path)
+ save_config(config, config_path)
+
+ # Save the model
+ if 'safetensors' in save_path:
+ safetensors.torch.save_file(state_dict, save_path, metadata={"format": "pt"})
+ else:
+ torch.save(state_dict, save_path)