More fixes.

Author: zxybazh

python/tvm/meta_schedule/cost_model/metric.py

@@ -18,7 +18,7 @@
 import numpy as np
 
 
-def max_curve(trial_scores):
+def max_curve(trial_scores: np.ndarray):
     """f(n) = max([s[i] fo i < n])
 
     Parameters
@@ -37,18 +37,3 @@ def max_curve(trial_scores):
         keep = max(keep, score)
         ret[i] = keep
     return ret
-
-
-def make_metric_sorter(focused_metric):
-    """ Make sure the focused metric is the first one. """
-
-    def metric_name_for_sort(name):
-        if focused_metric == name:
-            return "!" + name
-        return name
-
-    def sort_key(key):
-        key, _ = key
-        return metric_name_for_sort(key)
-
-    return sort_key

python/tvm/meta_schedule/cost_model/xgb_model.py

@@ -17,7 +17,7 @@
 """
 XGBoost-based cost model
 """
-from typing import Union, Optional, Tuple, Callable, List, TYPE_CHECKING
+from typing import Optional, Tuple, Callable, List, TYPE_CHECKING
 
 
 import logging
@@ -29,7 +29,7 @@
 from ..feature_extractor import FeatureExtractor
 from ..cost_model import PyCostModel
 from ..utils import cpu_count
-from .metric import max_curve, make_metric_sorter
+from .metric import max_curve
 
 if TYPE_CHECKING:
     from ..tune_context import TuneContext
@@ -39,6 +39,21 @@
 logger = logging.getLogger(__name__)
 
 
+def make_metric_sorter(focused_metric):
+    """ Make sure the focused metric is the first one. """
+
+    def metric_name_for_sort(name):
+        if focused_metric == name:
+            return "!" + name
+        return name
+
+    def sort_key(key):
+        key, _ = key
+        return metric_name_for_sort(key)
+
+    return sort_key
+
+
 class PackSum:
     """The pack-sum format
 
@@ -185,14 +200,9 @@ class XGBModel(PyCostModel):
         XGBoost model param, the eta, learning rate.
     xgb_seed : int
         XGBoost model param, the random seed.
-    xgb_nthread : int
+    xgb_nthread : Optional[int],
         XGBoost model param, the number of threads to use.
-    xgb_n_gpus : int
-        XGBoost model param, the number of gpus.
-    xgb_verbosity_train : int
-        XGBoost model param, the verbose level for training.
-    xgb_disable_default_eval_metric : Union[int, bool]
-        XGBoost model param, flag to disable default metric. Set to 1 or true to disable.
+        Default is None, which means to use physical number of cores.
     path : Optional[str]
         The path to save the model.
     num_warmup_samples : int
@@ -208,6 +218,23 @@ class XGBModel(PyCostModel):
 
     # model-related params
     _xgb_params: dict
+    """The parameters for xgboost model
+
+    Parameters
+    ----------
+    max_depth : int
+        XGBoost model param, the maximum depth.
+    gamma : float
+        XGBoost model param, the gamma.
+    min_child_weight : float
+        XGBoost model param, the minimum child weight.
+    eta : float
+        XGBoost model param, the eta, learning rate.
+    seed : int
+        XGBoost model param, the random seed.
+    nthread : int
+        XGBoost model param, the number of threads to use.
+    """
     # serialization-related
     path: Optional[str]
     # feature extractor
@@ -235,10 +262,7 @@ def __init__(
         xgb_min_child_weight: float = 0,
         xgb_eta: float = 0.2,
         xgb_seed: int = 43,
-        xgb_nthread: int = cpu_count(),
-        xgb_n_gpus: int = 0,
-        xgb_verbosity_train: int = 0,
-        xgb_disable_default_eval_metric: Union[int, bool] = 1,
+        xgb_nthread: Optional[int] = None,
         # load from disk
         path: Optional[str] = None,
         # behavior of randomness
@@ -252,16 +276,16 @@ def __init__(
         # feature extractor
         self.extractor = extractor
         # model-related
+        if xgb_nthread is None:
+            # use physical core number
+            xgb_nthread = cpu_count(False)
         self._xgb_params = {
             "max_depth": xgb_max_depth,
             "gamma": xgb_gamma,
             "min_child_weight": xgb_min_child_weight,
             "eta": xgb_eta,
             "seed": xgb_seed,
             "nthread": xgb_nthread,
-            "n_gpus": xgb_n_gpus,
-            "verbosity": xgb_verbosity_train,
-            "disable_default_eval_metric": xgb_disable_default_eval_metric,
         }
         # serialization-related
         self.path = path
@@ -277,7 +301,7 @@ def __init__(
         self.cached_normalizer = None
         self.booster = None
 
-    def load(self, path: str) -> None:
+    def load(self, path: str = None) -> None:
         """Load the cost model from given file location.
 
         Parameters
@@ -290,14 +314,18 @@ def load(self, path: str) -> None:
         Since XGBoost model trains from scratch, each time we can only load the model without the
         previous cached features / results so any call of update won't use previous training data.
         """
+        if path is None:
+            path = self.path
+
         import xgboost as xgb  # pylint: disable=import-outside-toplevel
 
         if self.booster is None:
             # save all the paramaters
             self.booster = xgb.Booster(self._xgb_params)
+        # throw error when path is none
         self.booster.load_model(path)
 
-    def save(self, path: str) -> None:
+    def save(self, path: str = None) -> None:
         """Save the cost model to given file location.
 
         Parameters
@@ -310,6 +338,9 @@ def save(self, path: str) -> None:
         Since XGBoost model trains from scratch, each time we can only save the model without the
         previous cached features / results so any call of update won't use previous training data.
         """
+        if path is None:
+            path = self.path
+        # throw error when path is none
         self.booster.save_model(path)
 
     def update(
@@ -322,7 +353,7 @@ def update(
 
         Parameters
         ----------
-        tune_context : TuneContext,
+        tune_context : TuneContext
             The tuning context.
         candidates : List[MeasureCandidate]
             The measure candidates.
@@ -339,7 +370,7 @@ def update(
             logger.debug(
                 "XGB validation: %s",
                 "\t".join(
-                    "%s: %.6f" % (key, score)  # pylint: disable=consider-using-f-string
+                    f"{key}: {score:.6f}"
                     for key, score in self._validate(
                         xs=new_features,
                         ys=new_mean_costs,
@@ -356,14 +387,11 @@ def update(
             xs=self.cached_features,
             ys=self.cached_mean_costs,
         )
-        # Update the model file if it has been set
-        if self.path:
-            self.save(self.path)
 
     def predict(
         self, tune_context: "TuneContext", candidates: List[MeasureCandidate]
     ) -> np.ndarray:
-        """Update the cost model given running results.
+        """Predict the normalized score using the cost model.
 
         Parameters
         ----------
@@ -375,7 +403,7 @@ def predict(
         Return
         ------
         result : np.ndarray
-            The predicted running results.
+            The predicted normalized score.
         """
         n_measured = len(self.cached_features)
         if self.booster is not None and n_measured >= self.num_warmup_samples:
@@ -407,8 +435,8 @@ def rmse(ys_pred: np.ndarray, d_train: "xgb.DMatrix"):  # pylint: disable = unus
             return self.d_train.rmse(ys_pred)
 
         def average_peak_score(
-            ys_pred: np.ndarray, d_train: "xgb.DMatrix"
-        ):  # pylint: disable = unused-argument
+            ys_pred: np.ndarray, d_train: "xgb.DMatrix"  # pylint: disable = unused-argument
+        ):
             return self.d_train.average_peak_score(ys_pred, self.average_peak_n)
 
         self.booster = xgb.train(
@@ -429,6 +457,12 @@ def average_peak_score(
             ],
         )
 
+        del self.d_train
+
+        # Update the model file if it has been set
+        if self.path:
+            self.save(self.path)
+
     def _predict(  # pylint: disable=invalid-name
         self,
         xs: List[np.ndarray],
@@ -565,9 +599,8 @@ def callback(env: "xgb.core.CallbackEnv"):
         if verbose_eval and iteration % verbose_eval == 0:
             info = []
             for key, score in eval_result:
-                if "null" in key:
-                    continue
-                info.append("%s: %.6f" % (key, score))  # pylint: disable=consider-using-f-string
+                if "null" not in key:
+                    info.append(f"{key}: {score:.6f}")
             logger.debug("XGB iter %3d: %s", iteration, "\t".join(info))
 
         ##### Choose score and do early stopping #####
@@ -581,10 +614,8 @@ def callback(env: "xgb.core.CallbackEnv"):
         best_score = state["best_score"]
         best_iteration = state["best_iteration"]
         if score < best_score:
-            msg = "[%d] %s" % (  # pylint: disable=consider-using-f-string
-                env.iteration,
-                "\t".join([_fmt_metric(x) for x in eval_result]),
-            )
+            tab = "\t"  # to work with f-string
+            msg = f"[{env.iteration}] {tab.join([_fmt_metric(x) for x in eval_result])}"
             state["best_msg"] = msg
             state["best_score"] = score
             state["best_iteration"] = env.iteration

python/tvm/meta_schedule/utils.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Utilities for meta schedule"""
+from logging import log
 from typing import Any, Callable, List, Optional, Union
 
 import ctypes