def test_hoisting_op_conv(): dtype = "float32" dshape = (1, 80, 73, 73) kshape = (192, 80, 3, 3) padding = (1, 1) groups = 1 dilation = (1, 1) kernel_size = (3, 3) channels = 192 scale = 1 x = relay.var("x", shape=dshape, dtype=dtype) w = relay.var("w", shape=kshape, dtype=dtype) y = relay.nn.conv2d( x, w, padding=padding, dilation=dilation, groups=groups, channels=channels, kernel_size=kernel_size, ) func = relay.Function([x, w], y) mod = tvm.IRModule() mod["main"] = func mod = relay.transform.InferType()(mod) data = np.random.uniform(-scale, scale, size=dshape).astype(dtype) kernel = np.random.uniform(-scale, scale, size=kshape).astype(dtype) params = {"w": tvm.nd.array(kernel)} for target, dev in enabled_targets(): with tvm.transform.PassContext(opt_level=3): lib = relay.build_module.build(mod, target=target, params=params) m = tvm.contrib.graph_executor.GraphModule(lib["default"](dev)) x = np.random.uniform(size=dshape) data_tvm = tvm.nd.array(data) m.set_input("x", data_tvm) m.run() e = m.module.time_evaluator("run", dev, number=300, repeat=3) t1 = e(data_tvm).results t1 = np.array(t1) * 1000 print("{} ms".format(t1.mean())) with tvm.transform.PassContext( opt_level=3, config={"tir.HoistIfThenElse": {"support_block_scope_hosting": True}} ): lib = relay.build_module.build(mod, target=target, params=params) m = tvm.contrib.graph_executor.GraphModule(lib["default"](dev)) x = np.random.uniform(size=dshape) data_tvm = tvm.nd.array(data) m.set_input("x", data_tvm) m.set_input(**params) m.run() e = m.module.time_evaluator("run", dev, number=300, repeat=3) t2 = e(data_tvm).results t2 = np.array(t2) * 1000 print("{} ms".format(t2.mean())) tvm.testing.assert_allclose(t1.mean(), t2.mean(), atol=1, rtol=1e-1)
def check_grad( func, inputs=None, test_inputs=None, eps=1e-6, atol=1e-5, rtol=1e-3, scale=None, mean=0, mode="higher_order", ): """Perform numerical gradient checking given a relay function. Compare analytical gradients to numerical gradients derived from two-sided approximation. Note that this test may fail if your function input types are not of high enough precision. Parameters ---------- func : tvm.relay.Function The relay function to test. inputs: List[np.array] Optional user-provided input parameters to use. If not given, will generate random normal inputs scaled to be close to the chosen epsilon value to avoid numerical precision loss. test_inputs: List[np.array] The inputs to test for gradient matching. Useful in cases where some inputs are not differentiable, such as symbolic inputs to dynamic ops. If not given, all inputs are tested. eps: float The epsilon value to use for computing numerical gradient approximation. atol: float The absolute tolerance on difference between numerical and analytical gradients. Note that this needs to be scaled appropriately relative to the chosen eps and inputs. rtol: float The relative tolerance on difference between numerical and analytical gradients. Note that this needs to be scaled appropriately relative to the chosen eps. scale: float The standard deviation of the inputs. mean: float The mean of the inputs. """ fwd_func = run_infer_type(func) bwd_func = run_infer_type(gradient(fwd_func, mode=mode)) if scale is None: scale = 10 * eps if inputs is None: params = fwd_func.params # Generate random inputs on the same scale as epsilon to avoid numerical precision loss. inputs = [_np_randn_from_type(x.checked_type, scale=scale, mean=mean) for x in params] if test_inputs is None: test_inputs = inputs for target, dev in enabled_targets(): # Eval the backward and forward functions # TODO(mbs): Evaluate a pair of functions so can share preparation between them. bwd_func_compiled = relay.create_executor(device=dev, target=target).evaluate(bwd_func) fwd_func_compiled = relay.create_executor(device=dev, target=target).evaluate(fwd_func) # Get analytic gradients. _, grads = bwd_func_compiled(*inputs) grads = [grad.numpy().astype("float64") for grad in grads] # Throw out gradients we aren't testing if inputs != test_inputs: tmp = [] # find the gradient that corresponds to every test input for test_input in test_inputs: for i, grad in enumerate(grads): if inputs[i] is test_input: tmp.append(grad) break grads = tmp assert len(grads) > 0, "You must test at least one gradient." # Get numeric gradients for each dimension of each param, using two-sided approximation. approx_grads = [] for x in test_inputs: approx_grad = np.zeros(x.shape) for i in np.ndindex(*x.shape): x_i = x[i] x[i] = x_i + eps fwd_plus = fwd_func_compiled(*inputs).numpy().astype("float64") x[i] = x_i - eps fwd_minus = fwd_func_compiled(*inputs).numpy().astype("float64") x[i] = x_i approx_grad[i] = np.sum((fwd_plus - fwd_minus) / (2 * eps)) approx_grads.append(approx_grad) # Compare gradients by checking that relative difference is below tolerance. for grad, approx_grad in zip(grads, approx_grads): np.testing.assert_allclose(grad, approx_grad, atol=atol, rtol=rtol)