[quant][graphmode] Swap dequantize after inline for ops that doesn't require observation (#33173)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/33173
How to deal with ops that’s defined for both floating point and quantized Tensor?
Category of ops: the ones that doesn’t require observers, which means the quantization parameters(scale/zero_point) of the output of this op can be inferred from the quantization parameters of inputs.
For example:
avg_pool, max_pool, flatten, transpose, upsample
Another related topic to previous one is how do we deal with things like adaptive_avg_pool2d that does not require to be observed and it works with quantized tensor as well? If we insert quant/dequant for them, even the quant fusion becomes a numerically changing operation because the scale/zero_point for input and output are different.
Proposal
We can swap the operator with dequantize whenever we see it. For example, for pattern
Let’s say aten::general_op is defined for both floating point and quantized
%r = aten::conv(...)
%q = quantize(%r)
%dq = dequantize(%q)
%f = aten::general_op(%dq)
...
We detect that all inputs of aten::general_op is produced by dequantize, we’ll first delete all the dequantize for the inputs and then insert dequantize for each use of the output of the aten::general_op, note that this should work generally for all the case we might encounter.
After transformation we’ll have:
%r = aten::conv(...)
%q = quantize(%r)
%x = aten::general_op(%q)
%f = dequantize(%x)
...
1. Multiple inputs
1. We need to make sure all inputs of the aten::general_op are produced by dequantize before we do this transformation
2. Input used by multiple operators
1. We already did this by inserting dequantize for each use of the value
3. Output used by multiple operators
1. We’ll reuse the code that inserts dequantize(might need some refactor)
Note that current concat does not belong to this category right now since it does not inherit quantization parameters from inputs.
Test Plan:
python test/test_jit.py
Imported from OSS
Differential Revision: D20123590
fbshipit-source-id: de2febe1f37e4079457a23acaeccbc6d9c9e1f8a
