tune_larq_icelake.py

import tvm
from tvm import relay, autotvm
import larq_zoo
import numpy as np
# import tvm.contrib.graph_runtime as runtime

import tvm.contrib.debugger.debug_runtime as runtime
import tvm.relay.testing
from tvm.autotvm.tuner import RandomTuner


batch_size = 1
dtype = "float32"
input_name = "input_1"
ctx = tvm.cpu()
target = "llvm -mcpu=icelake-client"
# target = "llvm -mcpu=cascadelake"

# data_layout = "NCHW"
data_layout = "NHWC"
ishape = [1, 3, 224, 224]

if data_layout == 'NHWC':
    ishape = [1, 224, 224, 3]


def tune_kernels(
    tasks,
    measure_option,
    n_trial=1000,
    tuner="gridsearch",
    early_stopping=None,
    log_filename="tuning.log",
):

    for i, task in enumerate(tasks):
        prefix = "[Task %2d/%2d] " % (i + 1, len(tasks))

        tuner_obj = RandomTuner(task)

        # do tuning
        # n_trial = len(task.config_space)
        tuner_obj.tune(
            n_trial=n_trial,
            early_stopping=early_stopping,
            measure_option=measure_option,
            callbacks=[
                autotvm.callback.progress_bar(n_trial, prefix=prefix),
                autotvm.callback.log_to_file(log_filename),
            ],
        )


def tune_and_evaluate(mod, params, tuning_opt):
    # extract workloads from relay program
    print("Extract tasks...")
    data_shape = tuple(ishape)
    out_shape = (1, 1000)

    tasks = autotvm.task.extract_from_program(
        mod["main"],
        target=target,
        params=params,
        ops=(relay.op.get("nn.conv2d"), relay.op.get("nn.bitserial_conv2d")),
    )

    tune_kernels(tasks, **tuning_opt)


if True:
    model_name = "resnet-18_bit_serial_%s" % data_layout
    model = larq_zoo.literature.BinaryResNetE18()

    larq_convert_map = relay.frontend.get_larq_convert_map()

    shape_dict = {"input_1": ishape}
    mod, params = relay.frontend.from_keras(
        model, shape_dict, custom_convert_map=larq_convert_map,
        layout=data_layout
    )
    # print(relay.transform.InferType()(mod)['main'])

else:
    model_name = "resnet-18_fp32"
    assert data_layout == "NCHW"
    mod, params = relay.testing.resnet.get_workload(
       num_layers=18,
       batch_size=1,
       layout='NCHW',
       dtype='float32',
       image_shape=(3, 224, 224))

log_file = "%s.log" % model_name
tuning_option = {
    "log_filename": log_file,
    "tuner": "random",
    'n_trial': 2000,
    'early_stopping': 800,
    "measure_option": autotvm.measure_option(
        builder=autotvm.LocalBuilder(),
        runner=autotvm.LocalRunner(
            number=1, repeat=10, min_repeat_ms=0, enable_cpu_cache_flush=True
        ),
    ),
}

if True:
    tune_and_evaluate(mod, params, tuning_option)

with autotvm.apply_history_best(log_file):
    with tvm.transform.PassContext(opt_level=3):
        graph, lib, params = relay.build_module.build(mod, target=target, params=params)

    module = runtime.create(graph, lib, ctx)
    dtype = "float32"
    data = np.random.randn(*ishape).astype(dtype)
    module.set_input(0, data)
    module.set_input(**params)
    module.run()# evaluate
    print("Evaluate inference time cost...")
    ftimer = module.module.time_evaluator("run", ctx, number=100, repeat=3)
    prof_res = np.array(ftimer().results) * 1000  # convert to millisecond
    print(
        "Mean inference time (std dev): %.2f ms (%.2f ms)"
        % (np.mean(prof_res), np.std(prof_res))
    )