Batch Norm folding clarification

Questions
#1

Hello everyone,

I am wondering about TVM’s ability to fold batch norm layers. As a refenrece, I will use the equations as shown here.

I know that TVM will constant fold the values in (1) which are to be multiplied/added with the output of the bias addition. In other words, there will be no square root nor division during inference. What I am wondering is can TVM do (4), (5) and (6). In other words, fold the batch norm onto the weights and biases.

If so, what is the pass that does this?

Thanks a lot

#2

It is the pass “FoldScaleAxis”.

#3

Thanks for your answer.

Weirdly enough in my case I still don’t see it working. I created a sequential pipeline of passes as follows [InferType(),SimplifyInference(),FoldConstant(),FoldScaleAxis()]

Are there any known limitations when the pass doesn’t work? (If it helps I am working on NHWC layout)

#4

I think the root cause of your problem is that the opt_level of this pass is bigger than the one you set, so it is skiped. Try this.

passes = [
    InferType(),
    SimplifyInference(),
    FoldConstant(),
    FoldScaleAxis()
]
with tvm.transform.PassContext(opt_level=3):
    nn_mod = tvm.transform.Sequential(passes)(nn_mod)
#5

Hi,

I think my problem is even more primitive but took me long to find. I think there isn’t an implementation of the FoldScaleAxis transformation for dense layers. Or at least I dont find it in the source file for that transform.

Do you agree?

CC: @dmitriy-arm since I see you did of the most recent changes to the pass, what would be the easiest way to add such a functionality for dense layer?