BatchNorm Fold into Weights

Problem

Fold batch normalization scale into convolution weights.

For each output channel o multiply the flattened weights by gamma[o] / sqrt(var[o] + eps) and round to 4 decimal places.

case 1

in: [[[2,4],[1,3]],[1,2],[1,4],0]

expected: [[2,4],[1,3]]

case 2

in: [[[1],[1]],[2,0.5],[1,1],0.000001]

expected: [[2],[0.5]]

case 3

in: [[[3],[6]],[1,1],[4,9],0]

expected: [[1.5],[2]]

Starter codePython

def fold_conv_weight(weight, gamma, var, eps):
    # Your code here
    pass

3 test cases · run your code against them in the editor