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)