[Relay][Quantization] KL-divergence-based per-layer calibration

[vinx13]

• KL divergence algorithm ported from MXNet quantization
• CollectStats pass that collects input of each simulated_quantize in annotated graph into a tuple output
• Support floating-point scale
• max_scale as an alternative for power2_scale in weight quantization

Evaluation code

https://gist.github.com/vinx13/6f1eb1f9e2c0a8786149ee881bfcd6aa

What's left:

• I added QAnnotateKind.BIAS. I'm not sure whether it is necessary. Currently there are a few tricks in handling bias (nbit_bias, valid_range, ...). It would be good to find a better solution and avoid these tricks.
• In my evaluation script, I have to write quantization workflow by myself (optimize, annotate, calibrate, realize). Please also share your thought on the design of calibrate function. We need to decide how users can specify different ways of quantization (max, power2, KLD, ...)

Evaluation result on ImageNet:

max_scale for weights, KL divergence for activations:

resnet18_v1, 0.70642 / 0.89702
resnet50_v1, 0.73682 / 0.91664
resnet101_v1, 0.74484 / 0.9208
resnet18_v2, 0.70794 / 0.89832
resnet50_v2, 0.7691 / 0.93268
resnet101_v2, 0.78204 / 0.94124

power2 for weights, KL divergence restricted to power2 value for activations (use --eval-power2 option in my evaluation script):

resnet18_v1, 0.70332 / 0.89526
resnet50_v1, 0.73426 / 0.9146
resnet101_v1, 0.72434 / 0.91058
resnet18_v2, 0.70314 / 0.89618
resnet50_v2, 0.76486 / 0.93108
resnet101_v2, 0.78066 / 0.94002

These experiments are done under opt_level=2. When opt_level=3, FoldScaleAxis might generate some outliers in bias vector and cause significant accuracy drops. We should use different scales than taking the maximum for bias in this case.

cc @tqchen@ZihengJiang@eqy@ajtulloch@antinucleon@FrozenGene

[vinx13]

This one is ready. Please review and share your thoughts on calibration api design.

[ZihengJiang]

If just collect stats, ExprVisitor should be enough

[vinx13]

ExprMutator is needed actually. This mutator transform annotated expr to original expr by removing each simulate_quantize.
For example Relay program:

%1 = ..
%2 = simulate_quantize(%1)
%3 = op(%2)
%4 = simulate_quantize(%3)

We need to profile %1 and %3. But %3 takes %2 as input, we need to replace input of %3 with %1 (because in Annotate pass simulate_quantize in %2 is not in passthrough mode, we need to either remove it or rewrite it in passthrough mode)

[vinx13]

@ZihengJiang I was thinking that the other pr #3543 actually breaks this pass (because the result of this pass contains annotations and casts)

[ZihengJiang]

@vinx13 Why not collect stats before annotate?

[vinx13]

@ZihengJiang Annotations tell us which nodes should be profiled. If we want to collect stats before annotate, we need to repeat the code similar to annotate to decide which node should be quantized.

[ZihengJiang]

Okay, let's keep the current way. #3543 will not breaks this pass since annotation.cast_hint and annotation.stop_fusion will not change the running result. They are just annotation and you can view them as identity. One thing is, instead of detecting and jumping simulated_quantize inside of IRMutator, let's adding an option like simulated_quantize(kind=kIdentity) for eliminating the impact of simulated_quantize

[vinx13]

@ZihengJiang updated

[ZihengJiang]

@vinx13 Could you please address other comments?
We can change the calibrate API like you did for now. In long term, we should think about like calibrate(graph, mod, ctx, fcalibrate), where fcalibrate(sq_op, stats) is a callback function which can be provided by user.

[tqchen]

@ZihengJiang@vinx13 please followup on this and let us merge soon

[zhenhuaw-me]

I leave some comments, please ping me if I understand in a wrong way.

[tqchen]

@ZihengJiang@jackwish please https://docs.tvm.ai/contribute/code_review.html#approve-and-request-changes-explicitly

[zhenhuaw-me]

LGTM basically.