def get_conv2d(batch_size,
IC,
HW,
OC,
KHKW,
Pad,
Stride,
layout="NCHW",
dtype="float32"):
from tvm.relay.testing.init import create_workload
from tvm.relay.testing import layers
data_layout = layout
kernel_layout = "OIHW" if layout == "NCHW" else "HWIO"
data_shape = (batch_size, IC, HW, HW)
data = relay.var("data", shape=data_shape, dtype=dtype)
net = layers.conv2d(data=data,
channels=OC,
kernel_size=(KHKW, KHKW),
strides=(Stride, Stride),
padding=(Pad, Pad),
name="conv2d_profile",
data_layout=data_layout,
kernel_layout=kernel_layout)
return create_workload(net)
def get_conv2d_depthwise(conv_type,
batch_size,
IC,
HW,
OC,
KHKW,
Pad,
Stride,
layout="NCHW",
dtype="float32"):
from tvm.relay.testing.init import create_workload
from tvm.relay.testing import layers
data_layout = layout
kernel_layout = "OIHW" if layout == "NCHW" else "HWIO"
data_shape = (batch_size, IC, HW, HW)
data = relay.var("data", shape=data_shape, dtype=dtype)
if 'conv' in conv_type:
net = layers.conv2d(data=data,
channels=OC,
kernel_size=(KHKW, KHKW),
strides=(Stride, Stride),
padding=(Pad, Pad),
name="conv2d_profile",
data_layout=data_layout,
kernel_layout=kernel_layout)
elif 'depthwise' in conv_type:
print("build depthwise net")
net = layers.conv2d(data=data,
channels=OC,
groups=OC,
kernel_size=(KHKW, KHKW),
strides=(Stride, Stride),
padding=(Pad, Pad),
name="depthwise_profile",
data_layout=data_layout,
kernel_layout=kernel_layout)
else:
print("err : not correct conv type")
return create_workload(net)
def tvm_lenet(num_classes=10,
data_shape=(1, 1, 32, 32),
dtype='float32',
alpha=1.0,
is_shallow=False):
from tvm import relay
from tvm.relay.testing import layers
"""Function to construct a Lenet"""
data = relay.var("data", shape=data_shape, dtype=dtype)
conv1 = layers.conv2d(data=data,
channels=6,
kernel_size=(5, 5),
name='conv1')
conv1 = relay.tanh(conv1)
pool2 = relay.nn.avg_pool2d(conv1, pool_size=(2, 2), strides=(2, 2))
conv3 = layers.conv2d(data=pool2,
channels=16,
kernel_size=(5, 5),
name='conv3')
conv3 = relay.tanh(conv3)
pool4 = relay.nn.avg_pool2d(conv3, pool_size=(2, 2), strides=(2, 2))
conv5 = layers.conv2d(data=pool4,
channels=120,
kernel_size=(5, 5),
name='conv5')
conv5 = relay.tanh(conv5)
# Temp
flattened6 = relay.reshape(conv5, (1, 120))
# flattened6 = relay.nn.batch_flatten(conv5)
fcw7 = relay.var('fc7_weight', shape=(120, 84))
fcw7 = relay.transpose(fcw7)
fc7 = relay.nn.dense(data=flattened6, weight=fcw7, units=84)
fc7 = relay.tanh(fc7)
fcw8 = relay.var('fc6_weight', shape=(84, 10))
fcw8 = relay.transpose(fcw8)
fc8 = relay.nn.dense(data=fc7, weight=fcw8, units=10)
softmax = relay.nn.softmax(data=fc8)
fn = relay.Function(relay.analysis.free_vars(softmax), softmax)
return fn
def get_operator(data_shape,
out_channel,
kernel_size,
strides,
padding,
dtype="float32"):
data = relay.var("data", shape=data_shape, dtype=dtype)
body = layers.conv2d(data=data,
channels=out_channel,
kernel_size=kernel_size,
strides=strides,
padding=padding,
name="conv2d")
return relay.Function(relay.ir_pass.free_vars(body), body)
def lenet(num_classes=10,
data_shape=(1, 1, 32, 32),
dtype='float32',
alpha=1.0,
is_shallow=False):
"""Function to construct a MobileNet"""
data = relay.var("data", shape=data_shape, dtype=dtype)
conv1 = layers.conv2d(data=data,
channels=6,
kernel_size=(5, 5),
name='conv1')
conv1 = relay.nn.relu(conv1)
pool2 = relay.nn.avg_pool2d(conv1, pool_size=(2, 2), strides=(2, 2))
conv3 = layers.conv2d(data=pool2,
channels=16,
kernel_size=(5, 5),
name='conv3')
conv3 = relay.nn.relu(conv3)
pool4 = relay.nn.avg_pool2d(conv3, pool_size=(2, 2), strides=(2, 2))
flattened5 = relay.nn.batch_flatten(pool4)
fcw5 = relay.var('fc5_weight')
fc5 = relay.nn.dense(data=flattened5, weight=fcw5, units=120)
fc5 = relay.nn.relu(fc5)
fcw6 = relay.var('fc6_weight')
fc6 = relay.nn.dense(data=fc5, weight=fcw6, units=84)
fc6 = relay.nn.relu(fc6)
fcw7 = relay.var('fc7_weight')
fc7 = relay.nn.dense(data=fc6, weight=fcw7, units=num_classes)
fc7 = relay.nn.relu(fc7)
softmax = relay.nn.softmax(data=fc7)
fn = relay.Function(relay.analysis.free_vars(softmax), softmax)
return init.create_workload(fn)
def tvm_generic(N,
H,
W,
C,
kernel_size,
K,
stride=1,
padding=0,
dilation=1,
groups=1,
number=100,
dev=0,
timeout=4,
target="llvm",
trials=100):
data_shape = (N, C, H, W)
data = relay.var("data", shape=data_shape, dtype="float32")
kernel_size = (kernel_size, kernel_size)
stride = (stride, stride)
padding = (padding, padding)
body = layers.conv2d(data=data,
channels=K,
kernel_size=kernel_size,
strides=stride,
padding=padding,
name="conv2d")
op = relay.Function(relay.ir_pass.free_vars(body), body)
sym, params = create_workload(op)
tasks = autotvm.task.extract_from_program(op,
target=target,
params=params,
ops=(relay.op.nn.conv2d, ))
tuning_option = {
"log_filename":
"tvm_baseline_{}.log".format(
(N, C, H, W, K, kernel_size, stride, padding, dilation, groups)),
"tuner":
"xgb",
"early_stopping":
30,
"measure_option":
autotvm.measure_option(
builder=autotvm.LocalBuilder(timeout=timeout),
runner=autotvm.LocalRunner(number=number,
repeat=1,
timeout=timeout,
min_repeat_ms=150),
# runner=autotvm.RPCRunner(
# '1080ti', # change the device key to your key
# '0.0.0.0', 9190,
# number=20, repeat=3, timeout=4, min_repeat_ms=150)
),
}
log_filename = tuning_option["log_filename"]
tuner = tuning_option["tuner"]
early_stopping = tuning_option["early_stopping"]
measure_option = tuning_option["measure_option"]
# only support one task
assert len(tasks) == 1
for i, task in enumerate(tasks):
prefix = "[Task %2d/%2d] " % (i + 1, len(tasks))
# create tuner
if tuner == 'xgb' or tuner == 'xgb-rank':
tuner_obj = XGBTuner(task, loss_type='rank')
elif tuner == 'ga':
tuner_obj = GATuner(task, pop_size=100)
elif tuner == 'random':
tuner_obj = RandomTuner(task)
elif tuner == 'gridsearch':
tuner_obj = GridSearchTuner(task)
else:
raise ValueError("Invalid tuner: " + tuner)
# do tuning
n_trial = trials
length = len(task.config_space)
print("config space length=", length)
# tuner_obj.tune(n_trial=min(n_trial, length),
# early_stopping=early_stopping,
# measure_option=measure_option,
# callbacks=[
# autotvm.callback.progress_bar(n_trial, prefix=prefix),
# autotvm.callback.log_to_file(log_filename)])
if not os.path.exists(log_filename):
raise RuntimeError(
"the log file {} doesn't exists".format(log_filename))
with autotvm.apply_history_best(log_filename):
with relay.build_config(opt_level=3):
graph, lib, params = relay.build_module.build(op,
target=target,
params=params)
ctx = tvm.device(str(target), 0)
data_tvm = tvm.nd.array(
(np.random.uniform(size=data_shape)).astype("float32"))
module = runtime.create(graph, lib, ctx)
module.set_input("data", data_tvm)
module.set_input(**params)
# evaluate
ftimer = module.module.time_evaluator("run",
ctx,
number=number,
repeat=1)
prof_res = np.array(ftimer().results) * 1e3
return prof_res
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)