[Convert Layout Pass] How can I change the layout of conv2d & dw_conv2d separately?

#1

I use this code to change the nn.conv2d layout.

# Set layouts
    desired_layouts = {'nn.conv2d': ['NHWC', 'OHWI']}
    seq = relay.transform.ConvertLayout(desired_layouts)
    mod = seq(mod)

but all conv2d layout is OHWI, I hope conv2d kernel_layout is OHWI, dw_conv2d kernel_layout is HWOI.

nn.conv2d(%dnnl_44_i0, %dnnl_44_i1, padding=[0, 0, 0, 0], channels=512, kernel_size=[1, 1], data_layout=“NHWC”, kernel_layout=“OHWI”, out_dtype=“int32”) //…ok

nn.conv2d(%dnnl_46_i0, %dnnl_46_i1, strides=[2, 2], padding=[0, 0, 0, 0], groups=512, channels=512, kernel_size=[3, 3], data_layout=“NHWC”, kernel_layout=“OHWI”, out_dtype=“int32”) //I hope OHWI -> HWOI

How can I change the layout of conv2d & dw_conv2d separately?

Thanks.

#2

Hi @Ryan ,

I ran into this issue a while ago now and ended up finding a solution based on @anijain2305’s suggestion. See the suggestion here: https://github.com/apache/tvm/pull/6430#issuecomment-689703548. Essentially you can override the default implementation of register_convert_op_layout which the convert layout pass will use to decide the correct layout of a particular ‘nn.conv2d’ operator. You can define this prior to running the convert layout pass, here is a rough untested example:

@reg.register_convert_op_layout("nn.conv2d", level=11)
def convert_conv2d(attrs, inputs, tinfos, desired_layouts):
    data, weight = inputs
    new_attrs = dict(attrs)

    is_depthwise = is_depthwise_conv2d(data.shape, attrs['data_layout'], weight.shape, attrs['kernel_layout'], attrs['groups'])

    new_attrs['data_layout'] = 'NHWC'
    if is_depthwise:
        new_attrs['kernel_layout'] = 'HWOI'
    else:
        new_attrs['kernel_layout'] = 'OHWI'
       
    return relay.nn.conv2d(data, weight, **new_attrs)

desired_layouts = {'nn.conv2d': ['NHWC', 'OHWI']}
mod = relay.transform.ConvertLayout(desired_layouts)(mod)

Hope this helps!

1 Like
#3

Thank you very much

I will try it.

#4

This method is OK.

Thank you.