| import tvm |
| import topi |
| import topi.testing |
| import numpy as np |
| |
| |
| def test_dilate(): |
| target = 'llvm' |
| ctx = tvm.cpu(0) |
| |
| def _test_dilate(input_size, strides): |
| Input = tvm.placeholder((input_size)) |
| Output = topi.nn.dilate(Input, strides) |
| schedule = tvm.create_schedule(Output.op) |
| input_np = np.random.uniform(size=input_size).astype(Input.dtype) |
| output_np = topi.testing.dilate_python(input_np, strides) |
| input_tvm = tvm.nd.array(input_np, ctx=ctx) |
| output_size = topi.util.get_const_tuple(Output.shape) |
| output_tvm = tvm.nd.array(np.zeros(shape=output_size).astype(Output.dtype), ctx=ctx) |
| f = tvm.build(schedule, [Input, Output], target) |
| f(input_tvm, output_tvm) |
| tvm.testing.assert_allclose(output_tvm.asnumpy(), output_np, rtol=1e-5) |
| |
| _test_dilate((32,), (2,)) |
| _test_dilate((32,32), (2,2)) |
| _test_dilate((1,3,32,32), (1,1,1,1)) |
| _test_dilate((1,3,32,32), (2,2,2,2)) |
| _test_dilate((1,32,32,3,3), (1,1,1,1,1)) |
| _test_dilate((1,32,32,3,3), (2,2,2,2,2)) |
| _test_dilate((1,32,32,32,3,3), (1,1,1,2,2,2)) |
| _test_dilate((1,32,32,32,3,3), (2,2,2,1,1,1)) |
| |
| |
| if __name__ == "__main__": |
| test_dilate() |
