Python tvm.relay方法代碼示例

# 需要導入模塊: import tvm [as 別名]
# 或者: from tvm import relay [as 別名]
def build_graph_relay(args, ctx, onnx_model, inputs, input_names):
    import tvm.relay as relay

    shape_dict = {}
    dtype_dict = {}
    for name, value in zip(input_names, inputs.values()):
        shape_dict[name] = value.shape
        dtype_dict[name] = value.dtype

    mod, params = relay.frontend.from_onnx(onnx_model, shape=shape_dict)
    with relay.build_config(opt_level=args.opt_level):
        graph, lib, params = relay.build(mod, args.target, params=params)

        if args.dump_frontend:
            print(graph)

    graph_module = create_graph_module(args, graph, lib, ctx)
    set_inputs(graph_module, ctx, inputs, params)

    return graph_module 

# 需要導入模塊: import tvm [as 別名]
# 或者: from tvm import relay [as 別名]
def relay_max_pool2d_grad(c, img, psize, stride, pad, dil, ceil_mode, dout):
    assert psize.is_constant(tuple)
    assert stride.is_constant(tuple)
    assert pad.is_constant(tuple)
    assert dil.is_constant(tuple)
    assert ceil_mode.is_constant(bool)
    assert dil.value == (1, 1)

    return relay.nn.max_pool2d_grad(
        c.ref(dout),
        c.ref(img),
        psize.value,
        stride.value,
        pad.value,
        ceil_mode=ceil_mode.value,
    ) 

# 需要導入模塊: import tvm [as 別名]
# 或者: from tvm import relay [as 別名]
def torch2tvm_module(torch_module: torch.nn.Module, torch_inputs: Tuple[torch.Tensor, ...], target):
    torch_module.eval()
    input_names = []
    input_shapes = {}
    with torch.no_grad():
        for index, torch_input in enumerate(torch_inputs):
            name = "i" + str(index)
            input_names.append(name)
            input_shapes[name] = torch_input.shape
        buffer = io.BytesIO()
        torch.onnx.export(torch_module, torch_inputs, buffer, input_names=input_names, output_names=["o" + str(i) for i in range(len(torch_inputs))])
        outs = torch_module(*torch_inputs)
        buffer.seek(0, 0)
        onnx_model = onnx.load_model(buffer)
        relay_module, params = tvm.relay.frontend.from_onnx(onnx_model, shape=input_shapes)
    with tvm.relay.build_config(opt_level=3):
        graph, tvm_module, params = tvm.relay.build(relay_module, target, params=params)
    return graph, tvm_module, params 

# 需要導入模塊: import tvm [as 別名]
# 或者: from tvm import relay [as 別名]
def example():
    shape = (1, 64, 54, 54)
    c_data = np.empty(shape).astype("float32")
    c = relay.const(c_data)
    weight = relay.var('weight', shape=(64, 64, 3, 3))
    x = relay.var("x", relay.TensorType((1, 64, 56, 56), "float32"))
    conv = relay.nn.conv2d(x, weight)
    y = relay.add(c, c)
    y = relay.multiply(y, relay.const(2, "float32"))
    y = relay.add(conv, y)
    z = relay.add(y, c)
    z1 = relay.add(y, c)
    z2 = relay.add(z, z1)
    return relay.Function([x], z2)

###############################################################################
# Let us register layout alteration for a conv2d op so that we can apply the
# layout alteration pass on the example. How alter layout pass works is out
# the scope of this tutorial. 

# 需要導入模塊: import tvm [as 別名]
# 或者: from tvm import relay [as 別名]
def alter_conv2d(attrs, inputs, tinfos, out_type):
    data, weight = inputs
    new_attrs = dict(attrs)
    new_attrs['data_layout'] = 'NCHW16c'
    return relay.nn.conv2d(data, weight, **new_attrs)

###############################################################################
# Optimize the Program
# --------------------
# Now we would like to optimize the program. Relay features a host of
# optimizations. We will select some of them to apply on this example program.
#
# There are multiple ways to optimize a Relay program. Below we will provide
# examples for each of them.
#
# Manually Apply Optimization Passes
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

# Let's first create a relay Module which contains one or multiple Relay
# functions for optimization. 

# 需要導入模塊: import tvm [as 別名]
# 或者: from tvm import relay [as 別名]
def run_relay_graph(mod, params, shape_dict, target, ctx):
    with relay.build_config(opt_level=3):
        graph, lib, params = relay.build(mod, target=target, params=params)
    input_shape = shape_dict["input_1"]
    dummy_data = np.random.uniform(size=input_shape, low=0, high=input_shape[1]).astype(
        "int32"
    )

    m = graph_runtime.create(graph, lib, ctx)
    m.set_input(0, dummy_data)
    m.set_input(**params)
    m.run()
    tvm_output = m.get_output(0)

    ftimer = m.module.time_evaluator("run", ctx, repeat=5, number=5)
    prof_res = np.array(ftimer().results) * 1000
    print(
        "%-20s %-19s (%s)"
        % ("Runtime:", "%.2f ms" % np.mean(prof_res), "%.2f ms" % np.std(prof_res))
    )
    return tvm_output 

# 需要導入模塊: import tvm [as 別名]
# 或者: from tvm import relay [as 別名]
def random_sparse_bert_params(func, params, density, BS_R, BS_C):
    def deepcopy(param_dic):
        ret = {}
        for k, v in param_dic.items():
            ret[k] = tvm.nd.array(v.asnumpy())
        return ret

    new_params = deepcopy(params)
    dense_weight_names = relay.analysis.sparse_dense._search_dense_op_weight(func)
    for item in dense_weight_names:
        name = str(item)
        shape = new_params[name].shape
        if shape[0] % BS_R == 0 and shape[1] % BS_C == 0:
            new_w = random_bsr_matrix(shape[0], shape[1], BS_R, BS_C, density)
            new_params[name] = tvm.nd.array(new_w)
    return new_params 

# 需要導入模塊: import tvm [as 別名]
# 或者: from tvm import relay [as 別名]
def run(args):
    onnx_model = onnx.load_model(run_onnx_util.onnx_model_file(args.test_dir, args.model_file))
    ctx = tvm.gpu()

    input_names, output_names = run_onnx_util.onnx_input_output_names(
        os.path.join(args.test_dir, args.model_file))

    test_data_dir = os.path.join(args.test_dir, 'test_data_set_0')
    inputs, outputs = run_onnx_util.load_test_data(
        test_data_dir, input_names, output_names)

    inputs = dict(inputs)
    graph_module = None
    if args.frontend == 'nnvm':
        graph_module = build_graph_nnvm(args, ctx, onnx_model, inputs, input_names)
    elif args.frontend == 'relay':
        graph_module = build_graph_relay(args, ctx, onnx_model, inputs, input_names)
    else:
        raise RuntimeError('Invalid frontend: {}'.format(args.frontend))

    graph_module.run()

    for i, (name, expected) in enumerate(outputs):
        tvm_output = tvm.nd.empty(expected.shape, expected.dtype, ctx=ctx)
        actual = graph_module.get_output(i, tvm_output).asnumpy()
        np.testing.assert_allclose(expected, actual,
                                   rtol=1e-3, atol=1e-4), name
        print('%s: OK' % name)
    print('ALL OK')

    def compute():
        graph_module.run()
        cupy.cuda.device.Device().synchronize()

    return run_onnx_util.run_benchmark(compute, args.iterations) 

# 需要導入模塊: import tvm [as 別名]
# 或者: from tvm import relay [as 別名]
def get_args(args=None):
    parser = argparse.ArgumentParser(description='Run ONNX by TVM')
    parser.add_argument('test_dir')
    parser.add_argument('--frontend', type=str, default='relay')
    parser.add_argument('--dump_frontend', action='store_true')
    parser.add_argument('--target', type=str, default='cuda')
    parser.add_argument('--debug', '-g', action='store_true')
    parser.add_argument('--iterations', '-I', type=int, default=1)
    parser.add_argument('--opt_level', '-O', type=int, default=3)
    parser.add_argument('--autotvm_log', type=str)
    parser.add_argument('--model_file', default=None)
    return parser.parse_args(args=args) 

# 需要導入模塊: import tvm [as 別名]
# 或者: from tvm import relay [as 別名]
def get_operator(data_shape, out_channel, kernel_size, strides, padding, dtype="float32"):
    data = relay.var("data", shape=data_shape, dtype=dtype)
    body = layers.conv2d(data=data, channels=out_channel, kernel_size=kernel_size, strides=strides, padding=padding, name="conv2d")
    return relay.Function(relay.ir_pass.free_vars(body), body) 

# 需要導入模塊: import tvm [as 別名]
# 或者: from tvm import relay [as 別名]
def tune_and_evaluate(tuning_opt, number, tune=True):
    op, params, data_shape = get_workload(batch_size, image_shape, out_channel, kernel_size, strides, padding)
    tasks = autotvm.task.extract_from_program(op, target=target, params=params, ops=(relay.op.nn.conv2d,))
    log_file = tuning_opt["log_filename"]
    if tune:
        print("Tuning...")
        tune_kernels(tasks, **tuning_opt)

    if not os.path.exists(log_file):
        raise RuntimeError("the log file {} doesn't exists".format(log_file))
    with autotvm.apply_history_best(log_file):
        print("Compile...")
        with relay.build_config(opt_level=3):
            graph, lib, params = relay.build_module.build(op, target=target, params=params)

        ctx = tvm.context(str(target), 0)
        data_tvm = tvm.nd.array((np.random.uniform(size=data_shape)).astype(dtype))
        module = runtime.create(graph, lib, ctx)
        module.set_input("data", data_tvm)
        module.set_input(**params)

        # evaluate
        print("Evaluate inference time cost...")
        ftimer = module.module.time_evaluator("run", ctx, number=1, repeat=number)
        prof_res = np.array(ftimer().results) * 1e3
        print("Time cost is: ", np.mean(prof_res)) 

# 需要導入模塊: import tvm [as 別名]
# 或者: from tvm import relay [as 別名]
def get_operator(data_shape, weight_shape, dtype="float32"):
    data = relay.var("data", shape=data_shape, dtype=dtype)
    weight = relay.var("weight", shape=weight_shape, dtype=dtype)
    body = relay.nn.dense(data=data, weight=weight, name="dense")
    return relay.Function(relay.ir_pass.free_vars(body), body) 

# 需要導入模塊: import tvm [as 別名]
# 或者: from tvm import relay [as 別名]
def tune_and_evaluate(tuning_opt, number, tune=True):
    op, params, data_shape = get_workload(batch_size, image_shape, out_channel, kernel_size, strides, padding)
    print(tvm.autotvm.task.task.TASK_TABLE)
    tasks = autotvm.task.extract_from_program(op, target=target, params=params, ops=(relay.op.nn.dense,))
    log_file = tuning_opt["log_filename"]
    if tune:
        print("Tuning...")
        tune_kernels(tasks, **tuning_opt)

    if not os.path.exists(log_file):
        raise RuntimeError("the log file {} doesn't exists".format(log_file))
    with autotvm.apply_history_best(log_file):
        print("Compile...")
        with relay.build_config(opt_level=3):
            graph, lib, params = relay.build_module.build(op, target=target, params=params)

        ctx = tvm.cpu(np.random.randint(0, 20))
        data_tvm = tvm.nd.array((np.random.uniform(size=data_shape)).astype(dtype))
        module = runtime.create(graph, lib, ctx)
        module.set_input("data", data_tvm)
        module.set_input(**params)

        # evaluate
        print("Evaluate inference time cost...")
        ftimer = module.module.time_evaluator("run", ctx, number=number, repeat=1)
        prof_res = np.array(ftimer().results) * 1e3
        cost = np.mean(prof_res)
        print("Time cost is: ", cost)
    return cost 

# 需要導入模塊: import tvm [as 別名]
# 或者: from tvm import relay [as 別名]
def tune_and_evaluate(tuning_opt, number, tune=True):
    op, params, data_shape = get_workload(batch_size, image_shape, out_channel, kernel_size, strides, padding)
    tasks = autotvm.task.extract_from_program(op, target=target, params=params, ops=(relay.op.nn.conv2d,))
    log_file = tuning_opt["log_filename"]
    if tune:
        print("Tuning...")
        tune_kernels(tasks, **tuning_opt)

    if not os.path.exists(log_file):
        raise RuntimeError("the log file {} doesn't exists".format(log_file))
    with autotvm.apply_history_best(log_file) as f:
        print(f.best_by_targetkey)
        print("Compile...")
        with relay.build_config(opt_level=3):
            graph, lib, params = relay.build_module.build(op, target=target, params=params)

        ctx = tvm.cpu(np.random.randint(0, 20))
        data_tvm = tvm.nd.array((np.random.uniform(size=data_shape)).astype(dtype))
        module = runtime.create(graph, lib, ctx)
        module.set_input("data", data_tvm)
        module.set_input(**params)

        # evaluate
        print("Evaluate inference time cost...")
        ftimer = module.module.time_evaluator("run", ctx, number=number, repeat=1)
        prof_res = np.array(ftimer().results) * 1e3
        cost = np.mean(prof_res)
        print("Time cost is: ", cost)
    return cost 

# 需要導入模塊: import tvm [as 別名]
# 或者: from tvm import relay [as 別名]
def get_network(name, batch_size):
    """Get the symbol definition and random weight of a network"""
    input_shape = (batch_size, 3, 224, 224)
    output_shape = (batch_size, 1000)

    if "resnet" in name:
        n_layer = int(name.split('-')[1])
        net, params = relay.testing.resnet.get_workload(num_layers=n_layer, batch_size=batch_size, dtype=dtype)
    elif "vgg" in name:
        n_layer = int(name.split('-')[1])
        net, params = relay.testing.vgg.get_workload(num_layers=n_layer, batch_size=batch_size, dtype=dtype)
    elif name == 'mobilenet':
        net, params = relay.testing.mobilenet.get_workload(batch_size=batch_size, dtype=dtype)
    elif name == 'squeezenet_v1.1':
        net, params = relay.testing.squeezenet.get_workload(batch_size=batch_size, version='1.1', dtype=dtype)
    elif name == 'inception_v3':
        input_shape = (1, 3, 299, 299)
        net, params = relay.testing.inception_v3.get_workload(batch_size=batch_size, dtype=dtype)
    elif name == 'mxnet':
        # an example for mxnet model
        from mxnet.gluon.model_zoo.vision import get_model
        block = get_model('resnet18_v1', pretrained=True)
        net, params = relay.frontend.from_mxnet(block, shape={'data': input_shape}, dtype=dtype)
        net = relay.Function(net.params, relay.nn.softmax(net.body), None, net.type_params, net.attrs)
    else:
        raise ValueError("Unsupported network: " + name)

    return net, params, input_shape, output_shape

# Replace "llvm" with the correct target of your cpu.
# For example, for AWS EC2 c5 instance with Intel Xeon
# Platinum 8000 series, the target should be "llvm -mcpu=skylake-avx512".
# For AWS EC2 c4 instance with Intel Xeon E5-2666 v3, it should be
# "llvm -mcpu=core-avx2".