[MetaSchedule] Added a cost model

apache · Jul 11, 2022 · ee752d0 · ee752d0
1 parent ce8f6d1
commit ee752d0
Show file tree

Hide file tree

Showing 4 changed files with 1,017 additions and 2 deletions.
diff --git a/python/tvm/meta_schedule/cost_model/__init__.py b/python/tvm/meta_schedule/cost_model/__init__.py
@@ -20,3 +20,8 @@
 from .cost_model import CostModel, PyCostModel
 from .random_model import RandomModel
 from .xgb_model import XGBModel
+
+try:
+    from .mlp_model import MLPModel  # type: ignore
+except ImportError:
+    pass
diff --git a/python/tvm/meta_schedule/cost_model/cost_model.py b/python/tvm/meta_schedule/cost_model/cost_model.py
@@ -190,7 +190,7 @@ def update(
         raise NotImplementedError
 
     def predict(self, context: TuneContext, candidates: List[MeasureCandidate]) -> np.ndarray:
-        """Update the cost model given running results.
+        """Predict given the measure candidates.
 
         Parameters
         ----------