Quantizatition and pruned model

OValery16 · June 24, 2021, 8:39am

Hi TVM community,

I am facing the following problem: I have pruned a 2D model and now I want to use TVM quantization. Since int8 quantization takes advantage of the dp4a primitive, the workload should be divisible by ic_block_factor which is 4. However, my network is pruned and the channels are no longer divisible by 4, which results in an error

github.com

apache/tvm/blob/5fa1c6dae0903f4dc31d39d42fcf582190ac1a68/python/tvm/topi/cuda/conv2d_int8.py#L91


output : tvm.te.Tensor
    5-D with shape [batch, out_channel_chunk, out_height, out_width, out_channel_block]
"""
assert layout in ["NCHW", "NCHW4c"]
ic_block_factor = 4
oc_block_factor = 4


pre_computed = len(kernel.shape) == 6
if not pre_computed:
    batch, channels, height, width = get_const_tuple(data.shape)
    assert (
        channels % ic_block_factor == 0
    ), "Number of input channels should be multiple of {}".format(ic_block_factor)
    packed_data = te.compute(
        (batch, channels // ic_block_factor, height, width, ic_block_factor),
        lambda n, c, h, w, vc: data[n, c * ic_block_factor + vc, h, w],
        name="packed_data",
    )


    out_channels, in_channels, kernel_h, kernel_w = get_const_tuple(kernel.shape)
    assert out_channels % 4 == 0, "Number of output channels should be multiple of {}".format(

I would like to pad the input weight channel/kernel. However, I am not really familiar with how TVM implements its kernel.

Can you advise me an approach to solve my problem?

reku · June 24, 2021, 11:37am

You can add your own legalize strategy, refer to this PR:

But I’m not sure if you will really get speed up after padding. So you can also consider directly modifying the cuda strategy of your own conv2d_int8 so that it can be distributed to topi that does not require dp4a.

github.com

apache/tvm/blob/5fa1c6dae0903f4dc31d39d42fcf582190ac1a68/python/tvm/relay/op/strategy/cuda.py#L119






@schedule_lrn.register(["cuda", "gpu"])
def schedule_lrn_cuda(attrs, outs, target):
    """schedule LRN for cuda"""
    with target:
        return topi.cuda.schedule_lrn(outs)




@conv2d_strategy.register(["cuda", "gpu"])
def conv2d_strategy_cuda(attrs, inputs, out_type, target):
    """conv2d cuda strategy"""
    strategy = _op.OpStrategy()
    data, kernel = inputs
    stride_h, stride_w = attrs.get_int_tuple("strides")
    dilation_h, dilation_w = attrs.get_int_tuple("dilation")
    padding = attrs.get_int_tuple("padding")
    groups = attrs.groups
    layout = attrs.data_layout
    kernel_layout = attrs.kernel_layout
    if dilation_h < 1 or dilation_w < 1: