def run_batchnorm(src_dtype, dst_dtype, rtol=1e-6, atol=1e-6):
shape = (3, 32, 32)
t = relay.TensorType(shape, src_dtype)
x = relay.var("x", t)
bn = batch_norm_infer(data=x, epsilon=2e-5, scale=False, name="bn_x")
f = relay.Function(relay.analysis.free_vars(bn), bn)
x_data = np.random.rand(*shape).astype(t.dtype)
module = tvm.IRModule.from_expr(f)
zero_data = np.zeros((32), "float32")
compare(
module,
(x_data, zero_data, zero_data, zero_data, zero_data),
src_dtype,
dst_dtype,
rtol,
atol,
)
def yolo():
inp = relay.var("data", shape=(1, 3, 416, 416))
conv0 = layers.conv2d(name="conv0",
data=inp,
channels=16,
kernel_size=(3, 3),
strides=(1, 1),
padding=(1, 1))
b0 = relay.var("b0_bias")
bias0 = relay.nn.bias_add(conv0, b0)
bn0 = layers.batch_norm_infer(bias0, name="bn0")
a1 = relay.nn.leaky_relu(bn0, alpha=0.1)
p2 = relay.nn.max_pool2d(a1, pool_size=(2, 2), strides=(2, 2))
conv3 = layers.conv2d(name="conv3",
data=p2,
channels=32,
kernel_size=(3, 3),
strides=(1, 1),
padding=(1, 1))
b3 = relay.var("b3_bias")
bias3 = relay.nn.bias_add(conv3, b3)
bn3 = layers.batch_norm_infer(bias3, name="bn3")
a4 = relay.nn.leaky_relu(bn3, alpha=0.1)
p5 = relay.nn.max_pool2d(a4, pool_size=(2, 2), strides=(2, 2))
conv6 = layers.conv2d(name="conv6",
data=p5,
channels=64,
kernel_size=(3, 3),
strides=(1, 1),
padding=(1, 1))
b6 = relay.var("b6_bias")
bias6 = relay.nn.bias_add(conv6, b6)
bn6 = layers.batch_norm_infer(bias6, name="bn6")
a7 = relay.nn.leaky_relu(bn6, alpha=0.1)
p8 = relay.nn.max_pool2d(a7, pool_size=(2, 2), strides=(2, 2))
conv9 = layers.conv2d(name="conv9",
data=p8,
channels=128,
kernel_size=(3, 3),
strides=(1, 1),
padding=(1, 1))
b9 = relay.var("b9_bias")
bias9 = relay.nn.bias_add(conv9, b9)
bn9 = layers.batch_norm_infer(bias9, name="bn9")
a10 = relay.nn.leaky_relu(bn9, alpha=0.1)
p11 = relay.nn.max_pool2d(a10, pool_size=(2, 2), strides=(2, 2))
conv12 = layers.conv2d(name="conv12",
data=p11,
channels=256,
kernel_size=(3, 3),
strides=(1, 1),
padding=(1, 1))
b12 = relay.var("b12_bias")
bias12 = relay.nn.bias_add(conv12, b12)
bn12 = layers.batch_norm_infer(bias12, name="bn12")
a13 = relay.nn.leaky_relu(bn12, alpha=0.1)
p14 = relay.nn.max_pool2d(a13, pool_size=(2, 2), strides=(2, 2))
conv15 = layers.conv2d(name="conv15",
data=p14,
channels=512,
kernel_size=(3, 3),
strides=(1, 1),
padding=(1, 1))
b15 = relay.var("b15_bias")
bias15 = relay.nn.bias_add(conv15, b15)
bn15 = layers.batch_norm_infer(bias15, name="bn15")
a16 = relay.nn.leaky_relu(bn15, alpha=0.1)
p17 = relay.nn.max_pool2d(a16,
pool_size=(2, 2),
strides=(1, 1),
padding=(0, 0))
conv18 = layers.conv2d(name="conv18",
data=p17,
channels=1024,
kernel_size=(3, 3),
strides=(1, 1),
padding=(1, 1))
b18 = relay.var("b18_bias")
bias18 = relay.nn.bias_add(conv18, b18)
bn18 = layers.batch_norm_infer(bias18, name="bn18")
a19 = relay.nn.leaky_relu(bn18, alpha=0.1)
conv20 = layers.conv2d(name="conv20",
data=a19,
channels=1024,
kernel_size=(3, 3),
strides=(1, 1),
padding=(1, 1))
b20 = relay.var("b20_bias")
bias20 = relay.nn.bias_add(conv20, b20)
bn20 = layers.batch_norm_infer(bias20, name="bn20")
a21 = relay.nn.leaky_relu(bn20, alpha=0.1)
conv22 = layers.conv2d(name="conv22",
data=a21,
channels=125,
kernel_size=(1, 1),
strides=(1, 1),
padding=(1, 1))
b22 = relay.var("b22_bias")
bias22 = relay.nn.bias_add(conv22, b22)
final = relay.op.add(bias22, relay.const(1.0))
fn = relay.Function(relay.analysis.free_vars(final), final)
return init.create_workload(fn)
