def tune_network(network, target):
# Extract tasks
mod, params = get_network(network)
target = tvm.target.Target(target)
tasks, task_weights = auto_scheduler.extract_tasks(mod["main"], params,
target)
with tempfile.NamedTemporaryFile() as fp:
log_file = fp.name
# Tuning
measure_ctx = auto_scheduler.LocalRPCMeasureContext(timeout=60)
tuner = auto_scheduler.TaskScheduler(tasks, task_weights)
tune_option = auto_scheduler.TuningOptions(
num_measure_trials=100,
num_measures_per_round=2,
early_stopping=1,
runner=measure_ctx.runner,
builder=auto_scheduler.LocalBuilder(timeout=60),
measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
)
tuner.tune(tune_option, search_policy="sketch.random")
del measure_ctx
# Compile with the history best
with auto_scheduler.ApplyHistoryBest(log_file):
with tvm.transform.PassContext(
opt_level=3,
config={"relay.backend.use_auto_scheduler": True}):
lib = relay.build(mod, target=target, params=params)
def test_tuning_cuda():
auto_scheduler.enable_relay_integration()
# Extract tasks
mod, params = get_network("mlp")
target = tvm.target.Target("cuda")
tasks, task_weights = auto_scheduler.extract_tasks(mod["main"], params, target)
objective = lambda costs: sum(c * w for c, w in zip(costs, task_weights))
with tempfile.NamedTemporaryFile() as fp:
log_file = fp.name
# Tuning
measure_ctx = auto_scheduler.LocalRPCMeasureContext(timeout=100)
tuner = auto_scheduler.TaskScheduler(tasks, objective)
tune_option = auto_scheduler.TuningOptions(
num_measure_trials=2,
num_measures_per_round=1,
runner=measure_ctx.runner,
measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
)
tuner.tune(tune_option, search_policy="sketch.random")
del measure_ctx
# Compile with the history best
with auto_scheduler.ApplyHistoryBest(log_file):
with tvm.transform.PassContext(opt_level=3):
lib = relay.build(mod, target=target, params=params)
# Todo(merrymercy): compile without any history to test the fallback mechanism
auto_scheduler.enable_relay_integration(False)
def tune_network(network, target):
# Extract tasks
mod, params = get_network(network)
target = tvm.target.Target(target)
tasks, task_weights = auto_scheduler.extract_tasks(mod["main"], params,
target)
with tempfile.NamedTemporaryFile() as fp:
log_file = fp.name
# Tuning
measure_ctx = auto_scheduler.LocalRPCMeasureContext(timeout=60)
tuner = auto_scheduler.TaskScheduler(tasks, task_weights)
tune_option = auto_scheduler.TuningOptions(
num_measure_trials=100,
num_measures_per_round=2,
early_stopping=1,
runner=measure_ctx.runner,
builder=auto_scheduler.LocalBuilder(timeout=60),
measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
)
tuner.tune(tune_option, search_policy="sketch.random")
del measure_ctx
# Compile with the history best
with auto_scheduler.ApplyHistoryBest(log_file):
with tvm.transform.PassContext(
opt_level=3,
config={"relay.backend.use_auto_scheduler": True}):
lib = relay.build(mod, target=target, params=params)
# Compile without auto-scheduler and any other optimization for correctness check
with tvm.transform.PassContext(opt_level=0):
lib2 = relay.build(mod, target=target, params=params)
# Check the correctness
def get_output(data, lib):
ctx = tvm.gpu()
module = graph_runtime.GraphModule(lib["default"](ctx))
module.set_input("data", data)
module.run()
return module.get_output(0).asnumpy()
np.random.seed(0)
if network == "mlp":
data = np.random.uniform(size=(1, 32))
elif network == "winograd-test":
data = np.random.uniform(size=(1, 23, 40, 32))
else:
raise ValueError("Unknown network: " + network)
actual_output = get_output(data, lib)
expected_output = get_output(data, lib2)
tvm.testing.assert_allclose(actual_output,
expected_output,
rtol=1e-4,
atol=1e-4)
def tune_network(network, target):
auto_scheduler.enable_relay_integration()
# Extract tasks
mod, params = get_network(network)
target = tvm.target.Target(target)
tasks, task_weights = auto_scheduler.extract_tasks(mod["main"], params,
target)
with tempfile.NamedTemporaryFile() as fp:
log_file = fp.name
# Tuning
measure_ctx = auto_scheduler.LocalRPCMeasureContext(timeout=60)
tuner = auto_scheduler.TaskScheduler(tasks, task_weights)
tune_option = auto_scheduler.TuningOptions(
num_measure_trials=100,
num_measures_per_round=2,
early_stopping=1,
runner=measure_ctx.runner,
builder=auto_scheduler.LocalBuilder(timeout=60),
measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
)
tuner.tune(tune_option, search_policy="sketch.random")
del measure_ctx
# Compile with the history best
with auto_scheduler.ApplyHistoryBest(log_file):
with tvm.transform.PassContext(opt_level=3):
lib = relay.build(mod, target=target, params=params)
# Todo(merrymercy): when the cpu backend is upstreamed, do the following things:
# 1. compile without history to test the fallback mechanism
# 2. check the correctness of layout rewrite / winograd pre-transform
auto_scheduler.enable_relay_integration(False)
def tune_network(network, target):
# Extract tasks
mod, params = get_network(network)
target = tvm.target.Target(target)
tasks, task_weights = auto_scheduler.extract_tasks(mod["main"], params, target)
with tempfile.NamedTemporaryFile() as fp:
log_file = fp.name
# Tuning
measure_ctx = auto_scheduler.LocalRPCMeasureContext(timeout=60, device=0)
tuner = auto_scheduler.TaskScheduler(tasks, task_weights, callbacks=[])
tune_option = auto_scheduler.TuningOptions(
num_measure_trials=100,
num_measures_per_round=2,
early_stopping=1,
runner=measure_ctx.runner,
builder=auto_scheduler.LocalBuilder(timeout=60),
measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
)
tuner.tune(tune_option, search_policy="sketch.random")
del measure_ctx
# Compile with the history best
with auto_scheduler.ApplyHistoryBest(log_file):
with tvm.transform.PassContext(
opt_level=3, config={"relay.backend.use_auto_scheduler": True}
):
lib = relay.build(mod, target=target, params=params)
# Also test that multiple log files can be loaded.
with auto_scheduler.ApplyHistoryBest([log_file, log_file]) as best:
assert isinstance(
best, auto_scheduler.dispatcher.ApplyHistoryBest
), "Unable to load multiple log files jointly."
# Confirm iterables can be directly loaded.
loaded_recs = auto_scheduler.dispatcher.load_records(log_file)
with auto_scheduler.ApplyHistoryBest(iter(loaded_recs)) as best:
assert isinstance(
best, auto_scheduler.dispatcher.ApplyHistoryBest
), "Unable to ingest logs from an interator."
# Sample a schedule when missing
with auto_scheduler.ApplyHistoryBestOrSample(None, num_measure=2):
with tvm.transform.PassContext(
opt_level=3, config={"relay.backend.use_auto_scheduler": True}
):
lib2 = relay.build(mod, target=target, params=params)
# Compile without auto-scheduler and any other optimization for correctness check
with tvm.transform.PassContext(opt_level=0):
ref_lib = relay.build(mod, target=target, params=params)
# Check the correctness
def get_output(data, lib):
dev = tvm.cuda()
module = graph_executor.GraphModule(lib["default"](dev))
module.set_input("data", data)
module.run()
return module.get_output(0).numpy()
np.random.seed(0)
if network == "mlp":
data = np.random.uniform(size=(1, 32))
elif network == "winograd-test":
data = np.random.uniform(size=(1, 23, 40, 32))
else:
raise ValueError("Unknown network: " + network)
actual_output1 = get_output(data, lib)
actual_output2 = get_output(data, lib2)
expected_output = get_output(data, ref_lib)
tvm.testing.assert_allclose(actual_output1, expected_output, rtol=1e-4, atol=1e-4)
tvm.testing.assert_allclose(actual_output2, expected_output, rtol=1e-4, atol=1e-4)
