def test_meta_schedule_cost_model():
@derived_object
class FancyCostModel(PyCostModel):
def load(self, path: str) -> None:
pass
def save(self, path: str) -> None:
pass
def update(
self,
context: TuneContext,
candidates: List[MeasureCandidate],
results: List[RunnerResult],
) -> None:
pass
def predict(self, context: TuneContext, candidates: List[MeasureCandidate]) -> np.ndarray:
return np.random.rand(10)
model = FancyCostModel()
model.save("fancy_test_location")
model.load("fancy_test_location")
model.update(TuneContext(), [], [])
results = model.predict(
TuneContext(), [MeasureCandidate(Schedule(mod=Matmul), []) for _ in range(10)]
)
assert results.shape == (10,)
def test_meta_schedule_random_model():
model = RandomModel()
model.update(TuneContext(), [], [])
res = model.predict(
TuneContext(),
[MeasureCandidate(Schedule(Matmul), []) for i in range(10)])
assert len(res) == 10
assert min(res) >= 0 and max(res) <= model.max_range
def test_meta_schedule_random_model_reseed():
model = RandomModel(seed=100)
res = model.predict(TuneContext(), [MeasureCandidate(Schedule(Matmul), []) for i in range(20)])
new_model = RandomModel(seed=100)
new_res = new_model.predict(
TuneContext(), [MeasureCandidate(Schedule(Matmul), []) for i in range(20)]
)
assert (res == new_res).all()
def test_meta_schedule_xgb_model():
extractor = RandomFeatureExtractor()
model = XGBModel(extractor=extractor, num_warmup_samples=2)
update_sample_count = 10
predict_sample_count = 100
model.update(
TuneContext(),
[_dummy_candidate() for i in range(update_sample_count)],
[_dummy_result() for i in range(update_sample_count)],
)
model.predict(TuneContext(), [_dummy_candidate() for i in range(predict_sample_count)])
def test_meta_schedule_random_model_reload():
model = RandomModel(seed=25973)
model.predict(
TuneContext(), [MeasureCandidate(Schedule(Matmul), []) for i in range(30)]
) # change state
path = os.path.join(tempfile.mkdtemp(), "test_output_meta_schedule_random_model.npy")
model.save(path)
res1 = model.predict(TuneContext(), [MeasureCandidate(Schedule(Matmul), []) for i in range(70)])
model.load(path)
res2 = model.predict(TuneContext(), [MeasureCandidate(Schedule(Matmul), []) for i in range(70)])
shutil.rmtree(os.path.dirname(path))
assert (res1 == res2).all()
def _create_context(mod, target, rule):
ctx = TuneContext(
mod=mod,
target=target,
space_generator=PostOrderApply(),
sch_rules=[rule],
task_name="test",
)
return ctx
def test_meta_schedule_design_space_generator_NIE():
@derived_object
class TestPySpaceGenerator(PySpaceGenerator):
pass
with pytest.raises(
TVMError, match="PySpaceGenerator's InitializeWithTuneContext method not implemented!"
):
generator = TestPySpaceGenerator()
generator.initialize_with_tune_context(TuneContext())
def _create_context(mod, target, rule):
ctx = TuneContext(
mod=mod,
target=target,
space_generator=PostOrderApply(),
sch_rules=[rule],
task_name="test",
)
ctx.space_generator.initialize_with_tune_context(ctx)
for sch_rule in ctx.sch_rules:
sch_rule.initialize_with_tune_context(ctx)
return ctx
def test_meta_schedule_xgb_model_reload():
extractor = RandomFeatureExtractor()
model = XGBModel(extractor=extractor, num_warmup_samples=10)
update_sample_count = 20
predict_sample_count = 30
model.update(
TuneContext(),
[_dummy_candidate() for i in range(update_sample_count)],
[_dummy_result() for i in range(update_sample_count)],
)
model.predict(TuneContext(),
[_dummy_candidate() for i in range(predict_sample_count)])
with tempfile.NamedTemporaryFile() as path:
# Backup
random_state = model.extractor.random_state # save feature extractor's random state
old_data = model.data
old_data_size = model.data_size
model.save(path.name)
res1 = model.predict(
TuneContext(),
[_dummy_candidate() for i in range(predict_sample_count)])
# Load
model.extractor.random_state = random_state # load feature extractor's random state
model.load(path.name)
new_data = model.data
new_data_size = model.data_size
res2 = model.predict(
TuneContext(),
[_dummy_candidate() for i in range(predict_sample_count)])
assert (res1 == res2).all()
assert old_data_size == new_data_size
assert len(old_data) == len(new_data)
for (k1, g1), (k2, g2) in zip(old_data.items(), new_data.items()):
assert k1 == k2
assert k1 == g1.group_hash
assert k2 == g2.group_hash
assert (g1.costs == g2.costs).all()
assert len(g1.features) == len(g2.features)
for f1, f2 in zip(g1.features, g2.features):
assert (f1 == f2).all()
def test_meta_schedule_xgb_model_reload():
extractor = RandomFeatureExtractor()
model = XGBModel(extractor=extractor, num_warmup_samples=10)
update_sample_count = 20
predict_sample_count = 30
model.update(
TuneContext(),
[_dummy_candidate() for i in range(update_sample_count)],
[_dummy_result() for i in range(update_sample_count)],
)
model.predict(TuneContext(),
[_dummy_candidate() for i in range(predict_sample_count)])
random_state = model.extractor.random_state # save feature extractor's random state
path = os.path.join(tempfile.mkdtemp(),
"test_output_meta_schedule_xgb_model.bin")
cached = (model.cached_features.copy(), model.cached_mean_costs.copy())
model.save(path)
res1 = model.predict(
TuneContext(),
[_dummy_candidate() for i in range(predict_sample_count)])
model.extractor.random_state = random_state # load feature extractor's random state
model.cached_features = None
model.cached_mean_costs = None
model.load(path)
new_cached = (model.cached_features.copy(), model.cached_mean_costs.copy())
res2 = model.predict(
TuneContext(),
[_dummy_candidate() for i in range(predict_sample_count)])
shutil.rmtree(os.path.dirname(path))
assert (res1 == res2).all()
# cached feature does not change
assert len(cached[0]) == len(new_cached[0])
for i in range(len(cached[0])):
assert (cached[0][i] == new_cached[0][i]).all()
# cached meaen cost does not change
assert (cached[1] == new_cached[1]).all()