65 lines
1.9 KiB
Python
65 lines
1.9 KiB
Python
import importlib
|
|
import numpy as np
|
|
import torch
|
|
import torch.distributed as dist
|
|
import os
|
|
|
|
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 check_istarget(name, para_list):
|
|
"""
|
|
name: full name of source para
|
|
para_list: partial name of target para
|
|
"""
|
|
istarget=False
|
|
for para in para_list:
|
|
if para in name:
|
|
return True
|
|
return istarget
|
|
|
|
|
|
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)
|
|
try:
|
|
obj = getattr(importlib.import_module(module, package=os.path.basename(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))), cls)
|
|
except:
|
|
obj = getattr(importlib.import_module(module, package=os.path.dirname(os.path.dirname(os.path.abspath( __file__ )))), cls)
|
|
return obj
|
|
|
|
|
|
def load_npz_from_dir(data_dir):
|
|
data = [np.load(os.path.join(data_dir, data_name))['arr_0'] for data_name in os.listdir(data_dir)]
|
|
data = np.concatenate(data, axis=0)
|
|
return data
|
|
|
|
def load_npz_from_paths(data_paths):
|
|
data = [np.load(data_path)['arr_0'] for data_path in data_paths]
|
|
data = np.concatenate(data, axis=0)
|
|
return data
|
|
|
|
def setup_dist(args):
|
|
if dist.is_initialized():
|
|
return
|
|
torch.cuda.set_device(args.local_rank)
|
|
torch.distributed.init_process_group(
|
|
'nccl',
|
|
init_method='env://'
|
|
) |