[MNT] modify the logic for qlib package init

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

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)

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
 

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

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

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=(",", ":"))

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

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)