I’m studying the VTA design and how it is being mapped to TVM. The resnet18 tutorial is good, however, the resnet18 itself is too complicated to follow. Instead, I’m trying with a simple nn.conv2d + nn.relu network as below:
def conv2d(data, weight=None, **kwargs):
name = kwargs.get("name")
kwargs.pop("name")
if not weight:
weight = relay.var(name + "_weight")
return relay.nn.conv2d(data, weight, **kwargs)
def conv_block(data, name, channels, kernel_size=(3, 3), strides=(1, 1),
padding=(1, 1), epsilon=1e-5):
conv = conv2d(
data=data,
channels=channels,
kernel_size=kernel_size,
strides=strides,
padding=padding,
data_layout='NCHW',
name=name+'_conv')
act = relay.nn.relu(data=conv)
return act
...
data_shape = (1, 3, 224, 224)
kernel_shape = (32, 3, 3, 3)
dtype = "float32"
data = relay.var("data", shape=data_shape, dtype=dtype)
act = conv_block(data, "graph", 32, strides=(2, 2))
func = relay.Function([data], act)
net = relay.frontend.common.infer_type(func)
...
mod = IRModule.from_expr(net)
...
with relay.build_config(opt_level=3, disabled_pass={"AlterOpLayout"}):
with vta.build_config(debug_flag=1):
graph, lib, params = relay.build(relay_prog, target, params=params, target_host=env.target_host
However, I noticed the conv2d is always being mapped to arm_cpu side, instead of VTA/FPGA side. Tracing the reason into vta_conv2d.py shows this is because of the data_layout=‘NCHW’ specified in conv2d definition. How can I change this definition to allow this conv2d to be mapped to VTA side? Thanks.
@_strategy.conv2d_strategy.register("vta")
def conv2d_strategy_vta(attrs, inputs, out_type, target):
"""conv2d vta strategy"""
strategy = OpStrategy()
kernel = inputs[1]
dilation = topi.util.get_const_tuple(attrs.dilation)
groups = attrs.groups
layout = attrs.data_layout
assert dilation == (1, 1), "support for dilation limited to (1, 1)"
if is_packed_layout(layout):
if groups == 1:
env = get_env()
assert env.LOG_INP_WIDTH == 3, "only support 8bit inp for now"
assert env.LOG_WGT_WIDTH == 3, "only support 8bit wgt for now"
assert kernel.dtype == "int8"
strategy.add_implementation(
_strategy.wrap_compute_conv2d(conv2d_packed, True),
_strategy.wrap_topi_schedule(schedule_conv2d_packed),
name="conv2d_packed.vta")
else: # group_conv2d
strategy.add_implementation(
_strategy.wrap_compute_conv2d(group_conv2d_packed, has_groups=True),
_strategy.wrap_topi_schedule(schedule_group_conv2d_packed),
name="group_conv2d_packed.vta")
return strategy
# If it's not packed, run on ARM CPU
print("vta/python/vta/top/op.py: conv2d layout not packed, run on ARM CPU.")
arm_tgt = tvm.target.arm_cpu(target.model)
return _strategy.arm_cpu.conv2d_strategy_arm_cpu(attrs, inputs, out_type, arm_tgt)