| """Test code for dense operator |
| Copied from topi/tests/python/test_topi_dense.py. |
| Should be removed once we fix OpenGL testing on Jenkins. |
| """ |
| import numpy as np |
| import tvm |
| import topi |
| from topi.util import get_const_tuple |
| from tvm.contrib.pickle_memoize import memoize |
| |
| |
| def verify_dense(batch, in_dim, out_dim, use_bias=True): |
| A = tvm.placeholder((batch, in_dim), name='A') |
| B = tvm.placeholder((out_dim, in_dim), name='B') |
| C = tvm.placeholder((out_dim,), name='C') |
| D = topi.nn.dense(A, B, C if use_bias else None) |
| D = topi.nn.relu(D) |
| dtype = A.dtype |
| |
| # use memoize to pickle the test data for next time use |
| @memoize("topi.tests.test_topi_dense") |
| def get_ref_data(): |
| a_np = np.random.uniform(size=(batch, in_dim)).astype(dtype) |
| b_np = np.random.uniform(size=(out_dim, in_dim)).astype(dtype) |
| c_np = np.random.uniform(size=(out_dim,)).astype(dtype) |
| if use_bias: |
| d_np = np.maximum(np.dot(a_np, b_np.T) + c_np, 0.0) |
| else: |
| d_np = np.maximum(np.dot(a_np, b_np.T), 0.0) |
| return (a_np, b_np, c_np, d_np) |
| # get the test data |
| a_np, b_np, c_np, d_np = get_ref_data() |
| |
| def check_device(device): |
| if not tvm.module.enabled(device): |
| print("Skip because %s is not enabled" % device) |
| return |
| print("Running on target: %s" % device) |
| with tvm.target.create(device): |
| s = topi.generic.schedule_dense(D) |
| ctx = tvm.context(device, 0) |
| a = tvm.nd.array(a_np, ctx) |
| b = tvm.nd.array(b_np, ctx) |
| c = tvm.nd.array(c_np, ctx) |
| d = tvm.nd.array(np.zeros(get_const_tuple(D.shape), dtype=dtype), ctx) |
| f = tvm.build(s, [A, B, C, D], device, name="dense") |
| f(a, b, c, d) |
| tvm.testing.assert_allclose(d.asnumpy(), d_np, rtol=1e-5) |
| |
| for device in ['opengl']: |
| check_device(device) |
| |
| def test_dense(): |
| verify_dense(1, 1024, 1000, use_bias=True) |
| verify_dense(1, 1024, 1000, use_bias=False) |
| |
| |
| if __name__ == "__main__": |
| test_dense() |
