diff --git a/learnware/__init__.py b/learnware/__init__.py index f081fc1..7445a56 100644 --- a/learnware/__init__.py +++ b/learnware/__init__.py @@ -2,31 +2,53 @@ __version__ = "0.1.2.99" import os from .logger import get_module_logger -from .utils import is_torch_available +from .utils import is_torch_available, setup_seed logger = get_module_logger("Initialization") -def init(make_dir: bool = False, tf_loglevel: str = "2", **kwargs): +def init(**kwargs): + """Init learnware package + + Parameters + ---------- + deterministic : bool, optional + whether to cancel randomness in learnware package, by default True + mkdir : bool, optional + whether to make directories for .learnware path, by default True + tf_loglevel: str, optional + The warning loglevel for tensforflow, by default "2" + """ from .config import C C.reset() - C.update(**kwargs) + C.update(**{k: v for k, v in kwargs.items() if k in C}) logger.info(f"init learnware market with {kwargs}") + + ## random seed + deterministic = kwargs.get("deterministic", True) + if deterministic: + setup_seed(C.random_seed) + ## make dirs - if make_dir: + mkdir = kwargs.get("mkdir", True) + if mkdir: os.makedirs(C.root_path, exist_ok=True) os.makedirs(C.database_path, exist_ok=True) + os.makedirs(C.stdout_path, exist_ok=True) + os.makedirs(C.cache_path, exist_ok=True) os.makedirs(C.learnware_pool_path, exist_ok=True) os.makedirs(C.learnware_zip_pool_path, exist_ok=True) os.makedirs(C.learnware_folder_pool_path, exist_ok=True) - logger.info(f"make learnware dir successfully!") ## ignore tensorflow warning - # os.environ["TF_CPP_MIN_LOG_LEVEL"] = tf_loglevel - # logger.info(f"The tensorflow log level is setted to {tf_loglevel}") + tf_loglevel = kwargs.get("tf_loglevel", "2") + os.environ["TF_CPP_MIN_LOG_LEVEL"] = tf_loglevel if not is_torch_available(verbose=False): logger.warning("The functionality of learnware is limited due to 'torch' is not installed!") + +# default init package +init() diff --git a/learnware/config.py b/learnware/config.py index 84c839b..d7b056f 100644 --- a/learnware/config.py +++ b/learnware/config.py @@ -65,11 +65,6 @@ DATABASE_PATH = os.path.join(ROOT_DIRPATH, "database") STDOUT_PATH = os.path.join(ROOT_DIRPATH, "stdout") CACHE_PATH = os.path.join(ROOT_DIRPATH, "cache") -# TODO: Delete them later -os.makedirs(ROOT_DIRPATH, exist_ok=True) -os.makedirs(DATABASE_PATH, exist_ok=True) -os.makedirs(STDOUT_PATH, exist_ok=True) -os.makedirs(CACHE_PATH, exist_ok=True) semantic_config = { "Data": { @@ -140,6 +135,7 @@ _DEFAULT_CONFIG = { "database_url": f"sqlite:///{DATABASE_PATH}", "max_reduced_set_size": 1310720, "backend_host": "http://www.lamda.nju.edu.cn/learnware/api", + "random_seed": 0, } C = Config(_DEFAULT_CONFIG) diff --git a/learnware/specification/base.py b/learnware/specification/base.py index 28212d6..064173d 100644 --- a/learnware/specification/base.py +++ b/learnware/specification/base.py @@ -19,6 +19,9 @@ class BaseStatSpecification: """Construct statistical specification""" raise NotImplementedError("generate_stat_spec_from_data is not implemented") + def get_states(self): + return {k: v for k, v in self.__dict__.items() if not k.startswith("_")} + def save(self, filepath: str): """Save the statistical specification into file in filepath diff --git a/learnware/specification/regular/image/rkme.py b/learnware/specification/regular/image/rkme.py index 50e367a..359ac89 100644 --- a/learnware/specification/regular/image/rkme.py +++ b/learnware/specification/regular/image/rkme.py @@ -40,7 +40,6 @@ class RKMEImageSpecification(RegularStatSpecification): self.beta = None self.cuda_idx = cuda_idx self.device = choose_device(cuda_idx=cuda_idx) - self.cache = False self.n_models = kwargs["n_models"] if "n_models" in kwargs else 16 self.model_config = ( @@ -371,7 +370,7 @@ class RKMEImageSpecification(RegularStatSpecification): rkme_load["z"] = torch.from_numpy(np.array(rkme_load["z"], dtype="float32")) rkme_load["beta"] = torch.from_numpy(np.array(rkme_load["beta"], dtype="float64")) - for d in self.__dir__(): + for d in self.__dict__(): if d in rkme_load.keys(): setattr(self, d, rkme_load[d]) diff --git a/learnware/specification/regular/table/rkme.py b/learnware/specification/regular/table/rkme.py index 8b97632..31410c7 100644 --- a/learnware/specification/regular/table/rkme.py +++ b/learnware/specification/regular/table/rkme.py @@ -444,7 +444,7 @@ class RKMETableSpecification(RegularStatSpecification): rkme_load["z"] = torch.from_numpy(np.array(rkme_load["z"])) rkme_load["beta"] = torch.from_numpy(np.array(rkme_load["beta"])) - for d in self.__dir__(): + for d in self.__dict__(): if d in rkme_load.keys(): setattr(self, d, rkme_load[d]) return True diff --git a/learnware/specification/system/hetero_table.py b/learnware/specification/system/hetero_table.py index 7e30710..e9086a8 100644 --- a/learnware/specification/system/hetero_table.py +++ b/learnware/specification/system/hetero_table.py @@ -10,13 +10,13 @@ import numpy as np from .base import SystemStatSpecification from ..regular import RKMETableSpecification from ..regular.table.rkme import torch_rbf_kernel -from ...utils import choose_device, setup_seed +from ...utils import choose_device, allocate_cuda_idx class HeteroMapTableSpecification(SystemStatSpecification): """Heterogeneous Map-Table Specification""" - def __init__(self, gamma: float = 0.1, cuda_idx: int = -1): + def __init__(self, gamma: float = 0.1, cuda_idx: int = None): """Initializing HeteroMapTableSpecification parameters. Parameters @@ -31,10 +31,9 @@ class HeteroMapTableSpecification(SystemStatSpecification): self.embedding = None self.weight = None self.gamma = gamma - self.cuda_idx = cuda_idx + self.cuda_idx = allocate_cuda_idx() if cuda_idx is None else cuda_idx torch.cuda.empty_cache() self.device = choose_device(cuda_idx=cuda_idx) - setup_seed(0) super(HeteroMapTableSpecification, self).__init__(type=self.__class__.__name__) def get_z(self) -> np.ndarray: @@ -125,11 +124,10 @@ class HeteroMapTableSpecification(SystemStatSpecification): with codecs.open(load_path, "r", encoding="utf-8") as fin: obj_text = fin.read() embedding_load = json.loads(obj_text) - embedding_load["device"] = choose_device(embedding_load["cuda_idx"]) - embedding_load["z"] = torch.from_numpy(np.array(embedding_load["z"])) + embedding_load["device"] = choose_device(["cuda_idx"]) + embedding_load["z"] = torch.from_numpy(np.arraembedding_loady(embedding_load["z"])) embedding_load["beta"] = torch.from_numpy(np.array(embedding_load["beta"])) - - for d in self.__dir__(): + for d in self.__dict__(): if d in embedding_load.keys(): setattr(self, d, embedding_load[d]) return True @@ -153,8 +151,5 @@ class HeteroMapTableSpecification(SystemStatSpecification): embedding_to_save["beta"] = embedding_to_save["beta"].detach().cpu().numpy() embedding_to_save["beta"] = embedding_to_save["beta"].tolist() embedding_to_save["device"] = "gpu" if embedding_to_save["cuda_idx"] != -1 else "cpu" - json.dump( - embedding_to_save, - codecs.open(save_path, "w", encoding="utf-8"), - separators=(",", ":"), - ) + with codecs.open(save_path, "w", encoding="utf-8") as fout: + json.dump(embedding_to_save, fout, separators=(",", ":")) diff --git a/learnware/utils/__init__.py b/learnware/utils/__init__.py index c118620..d98e60b 100644 --- a/learnware/utils/__init__.py +++ b/learnware/utils/__init__.py @@ -4,7 +4,7 @@ import zipfile from .import_utils import is_torch_available from .module import get_module_by_module_path from .file import read_yaml_to_dict, save_dict_to_yaml -from .gpu import setup_seed, choose_device +from .gpu import setup_seed, choose_device, allocate_cuda_idx def zip_learnware_folder(path: str, output_name: str): diff --git a/learnware/utils/gpu.py b/learnware/utils/gpu.py index b397221..cb856d4 100644 --- a/learnware/utils/gpu.py +++ b/learnware/utils/gpu.py @@ -47,37 +47,14 @@ def choose_device(cuda_idx=-1): return device -class CudaManager: - def __init__(self): - if is_torch_available(verbose=False): - import torch - - self.cuda_avalable = torch.cuda.is_available() - self.cuda_count = torch.cuda.device_count() if self.cuda_avalable else 0 - else: - self.cuda_avalable = False - self.cuda_count = 0 - - self.cur_cuda_idx = 0 - self.stat_spec_cuda = {} - - def reset(self): - self.cur_cuda_idx = 0 - self.stat_spec_cuda = {} - - def allocate_cuda(self): - if not self.cuda_avalable: - return -1 - - ret_cuda_idx = self.cur_cuda_idx - self.cur_cuda_idx = (self.cur_cuda_idx + 1) % self.cuda_count - return ret_cuda_idx - - def allocate_stat_spec_cuda(self, stat_spec): - if stat_spec.type not in self.stat_spec_cuda: - self.stat_spec_cuda[stat_spec.type] = self.allocate_cuda() - - return self.stat_spec_cuda[stat_spec.type] +def allocate_cuda_idx(self): + if is_torch_available(verbose=False): + import torch + cuda_count = torch.cuda.device_count() if torch.cuda.is_available() else 0 + else: + cuda_count = 0 -cuda_manager = CudaManager() + if cuda_count == 0: + return -1 + return np.random.randint(0, cuda_count)