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_dense.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) ``` --- [Visit Topic](https://discuss.tvm.ai/t/how-to-map-nn-conv2d-to-vta/6372/1) to respond. You are receiving this because you enabled mailing list mode. To unsubscribe from these emails, [click here](https://discuss.tvm.ai/email/unsubscribe/420d119a430a64d7cba61f6c1856253d14017c875452a07b991eda0b25ec6453).