test_build.py 5.43 KB
Newer Older
1 2 3
import numpy as np

import tvm
4
from tvm.contrib import graph_runtime
5 6
import nnvm.symbol as sym
import nnvm.compiler
7
from nnvm.compiler.build_module import _run_graph, precompute_prune
8 9 10 11 12 13 14 15

def test_compile():
    x = sym.Variable("x")
    y = sym.Variable("y")
    z = sym.exp(y + x)
    shape = (10, 128)
    dtype = tvm.float32
    shape_dict = {"x": shape, "y": shape}
16
    def verify(graph, lib):
17
        m = graph_runtime.create(graph, lib, tvm.cpu(0))
18 19 20 21 22 23 24 25 26 27 28 29
        # get member functions
        set_input, run, get_output = m["set_input"], m["run"], m["get_output"]
        na = tvm.nd.array(np.random.uniform(size=shape).astype(dtype))
        nb = tvm.nd.array(np.random.uniform(size=shape).astype(dtype))
        # set inputs
        set_input("x", na)
        set_input("y", nb)
        # execute
        run()
        # get outputs
        out = tvm.nd.empty(shape, dtype)
        get_output(0, out)
30
        tvm.testing.assert_allclose(
31 32 33 34 35 36 37 38 39 40 41 42
            out.asnumpy(), np.exp(na.asnumpy() + nb.asnumpy()))

    graph, lib, _ = nnvm.compiler.build(z, "llvm", shape_dict)
    assert graph.index.num_nodes == 3
    verify(graph, lib)

    with nnvm.compiler.build_config(opt_level=0):
        graph, lib, _ = nnvm.compiler.build(z, "llvm", shape_dict)
        # print(graph.ir())
        assert graph.index.num_nodes == 4
        verify(graph, lib)

43 44 45 46 47 48 49 50
def test_run():
    x = sym.Variable("x")
    y = sym.Variable("y")
    z = sym.exp(y + x)
    shape = (10, 10)
    dtype = tvm.float32
    nx = tvm.nd.array(np.random.uniform(size=shape).astype(dtype))
    ny = tvm.nd.array(np.random.uniform(size=shape).astype(dtype))
51
    res = _run_graph(z, {"x": nx, "y": ny})
52
    tvm.testing.assert_allclose(
53 54 55 56 57
        res[0].asnumpy(), np.exp(nx.asnumpy() + ny.asnumpy()))


def test_precompute_prune():
    x = sym.Variable("x") + 1
58
    a = sym.Variable("a")
59
    y = sym.Variable("y")
60
    z = y + x + a
61 62 63
    shape = (10, 10)
    dtype = tvm.float32
    nx = tvm.nd.array(np.random.uniform(size=shape).astype(dtype))
64
    na = tvm.nd.array(np.random.uniform(size=shape).astype(dtype))
65
    ny = tvm.nd.array(np.random.uniform(size=shape).astype(dtype))
66
    params = {"x": nx, "a": na}
67 68
    graph, lib, params = nnvm.compiler.build(
        z, "llvm", shape={"y": ny.shape}, params=params)
69
    assert graph.index.num_nodes == 4
70
    m = graph_runtime.create(graph, lib, tvm.cpu(0))
71 72
    params["y"] = ny
    res = tvm.nd.empty(shape)
73 74
    m["load_params"](nnvm.compiler.save_param_dict(params))
    m.run()
75
    out = m.get_output(0, out=res)
76
    tvm.testing.assert_allclose(
77
        res.asnumpy(), nx.asnumpy() + 1 + ny.asnumpy() + na.asnumpy())
78 79


80 81 82 83 84 85 86 87 88 89 90 91 92 93 94
def test_dtypes():
    x = sym.Variable("x")
    y = sym.relu(x)
    dshape = (1, 3, 32, 32)
    oshape = dshape
    for dtype in ['float32', 'float64', 'int32', 'int16', 'int8', 'int64']:
        graph, lib, _ = nnvm.compiler.build(y, 'llvm', {"x": dshape}, dtype=dtype)
        m = graph_runtime.create(graph, lib, tvm.cpu())
        if 'float' in dtype:
          data = np.random.uniform(size=dshape).astype(dtype)
        elif 'int' in dtype:
          data = np.random.randint(-127, 127, dshape).astype(dtype)
        m.run(x=data)
        data = (data > 0) * data
        out = m.get_output(0, tvm.nd.empty(oshape, dtype))
95
        tvm.testing.assert_allclose(out.asnumpy(), data, atol=1e-5, rtol=1e-5)
96

97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112
def test_ndarray_output():
    x = sym.Variable("x")
    y = sym.Variable("y")
    z = x + y
    shape = (10, 10)
    dtype = tvm.float32
    nx = tvm.nd.array(np.random.uniform(size=shape).astype(dtype))
    ny = tvm.nd.array(np.random.uniform(size=shape).astype(dtype))
    params = {"x": nx, "ny": ny}
    graph, lib, params = nnvm.compiler.build(
        z, "llvm", shape={"y": ny.shape, "x": nx.shape}, params=params)
    m = graph_runtime.create(graph, lib, tvm.cpu(0))
    m.set_input("x", nx)
    m.set_input("y", ny)
    m.run()
    out = m.get_output(0)
113
    tvm.testing.assert_allclose(
114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133
        out.asnumpy(), nx.asnumpy() + ny.asnumpy())

def test_ndarray_input():
    x = sym.Variable("x")
    y = sym.Variable("y")
    z = x + y
    shape = (10, 10)
    dtype = tvm.float32
    nx = tvm.nd.array(np.random.uniform(size=shape).astype(dtype))
    ny = tvm.nd.array(np.random.uniform(size=shape).astype(dtype))
    params = {"x": nx, "ny": ny}
    graph, lib, params = nnvm.compiler.build(
        z, "llvm", shape={"y": ny.shape, "x": nx.shape}, params=params)
    m = graph_runtime.create(graph, lib, tvm.cpu(0))
    m.set_input("x", nx)
    m.set_input("y", ny)
    in_x = tvm.nd.empty(shape, dtype)
    in_y = tvm.nd.empty(shape, dtype)
    m.get_input("x", in_x)
    m.get_input("y", in_y)
134 135
    tvm.testing.assert_allclose(nx.asnumpy(), in_x.asnumpy())
    tvm.testing.assert_allclose(ny.asnumpy(), in_y.asnumpy())
136 137
    in_nx = m.get_input("x")
    in_ny = m.get_input("y")
138 139
    tvm.testing.assert_allclose(nx.asnumpy(), in_nx.asnumpy())
    tvm.testing.assert_allclose(ny.asnumpy(), in_ny.asnumpy())
140 141 142 143 144 145 146 147 148 149 150 151

def test_num_outputs():
    x = sym.Variable('x')
    z = sym.split(x, indices_or_sections=5, axis=1)
    shape = (10, 10)
    dtype = tvm.float32
    nx = tvm.nd.array(np.random.uniform(size=shape).astype(dtype))
    params = {"x": nx}
    graph, lib, params = nnvm.compiler.build(
        z, "llvm", shape={"x": nx.shape}, params=params)
    m = graph_runtime.create(graph, lib, tvm.cpu(0))
    assert m.get_num_outputs() == 5
152

153
if __name__ == "__main__":
154
    test_precompute_prune()
155
    test_compile()
156
    test_run()
157
    test_dtypes()
158 159 160
    test_ndarray_output()
    test_ndarray_input()
    test_num_outputs()