[MetaSchedule] Added a cost model

apache · Jul 7, 2022 · c39249d · c39249d
1 parent ce8f6d1
commit c39249d
Show file tree

Hide file tree

Showing 3 changed files with 781 additions and 1 deletion.
diff --git a/python/tvm/meta_schedule/cost_model/__init__.py b/python/tvm/meta_schedule/cost_model/__init__.py
@@ -18,5 +18,6 @@
 The tvm.meta_schedule.cost_model package.
 """
 from .cost_model import CostModel, PyCostModel
+from .mlp_model import MLPModel
 from .random_model import RandomModel
 from .xgb_model import XGBModel
diff --git a/python/tvm/meta_schedule/cost_model/cost_model.py b/python/tvm/meta_schedule/cost_model/cost_model.py
@@ -175,6 +175,7 @@ def update(
         context: TuneContext,
         candidates: List[MeasureCandidate],
         results: List[RunnerResult],
+        frozen: bool = False,
     ) -> None:
         """Update the cost model given running results.
 
@@ -186,11 +187,13 @@ def update(
             The measure candidates.
         results : List[RunnerResult]
             The running results of the measure candidates.
+        frozen : bool
+            Skip updating the cost model.
         """
         raise NotImplementedError
 
     def predict(self, context: TuneContext, candidates: List[MeasureCandidate]) -> np.ndarray:
-        """Update the cost model given running results.
+        """Predict the normalized score using the cost model.
 
         Parameters
         ----------