[RELAY][PASS] FoldScaleAxis Backward

[tqchen]

This is a followup of #2020 , this PR implemented the infrastructure to do backward folding of scale on axis.

Goal

Fold the scaling of axis(usually caused by BatchNorm) into weight of conv2d in the past. For example

Old:

%1 = conv2d(%1, %w, data_layout="NHWC")
%2 = multiply(%x, %scale)

Transformed:

# scale weight's output channel
%1 = multiply(%w, expand_dims(%scale, axis=1, num_newaxis=3))
%2 = conv2d(%x, %1, data_layout="NHWC")

Further constant folding can fold the multiplication and we remove the scaling in the network.

The Algorithm

The general idea is similar to the Forward algorithm. We pass additional argument (axes, scale), when doing transformation, to request the result to satisfy

result = value
for i, k in enumerate(axes):
   k-ith dimension of result *= i-th dimension of scale

The problem is again we don;t want to blindly propagate scale backward it won't get consumed. So we run a forward "preparation phase", which propagates the demand of the potential axes scaling.

The new pass is more general than the FoldScaleAxis in nnvm

• The new pass support arbitrary scaling of multiple axes(although the further implementation is necessary), which could be helpful in NHWCc case,
• The new pass support folding backward into sum of two of conv2d, which was not possible in NNVM.

[tqchen]

@jroesch @ZihengJiang @yzhliu @zhiics @FrozenGene @merrymercy @srkreddy1238 @masahi please review

[FrozenGene]

Could we use is_const_int(wshape[0], param->groups) to check depthwise convolution? In theory, wshape[1] could be not equal to 1, although TF / MXNet / Caffe and so on frontends will make it be 1 even though the depth channel multiplier not be 1 (such as 0.25)

[tqchen]

The current code makes the assumption of input and output groups being even, and the special logic relies on this

A better approach would enhance the Conv2D code itself to handle general group convolution, which should not be too hard as a followup contribution from anyone in the community:)

[FrozenGene]

How about add clip operator like relu? Like Tensorflow frontend or Keras frontend, which use clip to implement relu6 operator.

[tqchen]

clip can be supported, this PR mainly introduces the scaffolding of things necessary, with new folder infrastructure, we can support clip as well as other operators(flatten, transpose) by more registrations

[masahi]

backward

[masahi]

for

[tqchen]

Thanks, @masahi @FrozenGene I have updated the PR according to your comments

[tqchen]

Thanks @FrozenGene @masahi , this is merged. Followup PRs to support more ops(clip) and grouped conv is more than welcomed!