Modify the program to expect quantized input at given index. The input is expected to be quantizing this input as the first step. Must be called before permute_input_layout. Returns the scale, zero point, qmin, qmax, and dtype of the expected quantization.
(
exported_program, input_index, qparams: Optional[Dict[str, Any]] = None
)
| 28 | |
| 29 | |
| 30 | def quantize_input( |
| 31 | exported_program, input_index, qparams: Optional[Dict[str, Any]] = None |
| 32 | ): |
| 33 | """ |
| 34 | Modify the program to expect quantized input at given index. The input is expected |
| 35 | to be quantizing this input as the first step. Must be called before |
| 36 | permute_input_layout. Returns the scale, zero point, qmin, qmax, and dtype of the |
| 37 | expected quantization. |
| 38 | """ |
| 39 | graph = exported_program.graph_module.graph |
| 40 | name = exported_program.graph_signature.user_inputs[input_index] |
| 41 | placeholders = [n for n in graph.nodes if n.op == "placeholder" and n.name == name] |
| 42 | assert placeholders |
| 43 | target_placeholder = placeholders[0] |
| 44 | |
| 45 | if len(target_placeholder.users) != 1: |
| 46 | raise ValueError(f"Input {input_index} has more than one users") |
| 47 | quantize = next(iter(target_placeholder.users)) |
| 48 | if quantize.target not in [ |
| 49 | exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, |
| 50 | torch.ops.quantized_decomposed.quantize_per_tensor.default, |
| 51 | ]: |
| 52 | raise ValueError( |
| 53 | f"Input {input_index} is not used by a quantize op. It's used by {quantize.target}" |
| 54 | ) |
| 55 | |
| 56 | if ( |
| 57 | quantize.target |
| 58 | == exir_ops.edge.quantized_decomposed.quantize_per_tensor.default |
| 59 | ): |
| 60 | replacement_op_dequant = ( |
| 61 | exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default |
| 62 | ) |
| 63 | replacement_op_quant = ( |
| 64 | exir_ops.edge.quantized_decomposed.quantize_per_tensor.default |
| 65 | ) |
| 66 | elif quantize.target == torch.ops.quantized_decomposed.quantize_per_tensor.default: |
| 67 | replacement_op_dequant = ( |
| 68 | torch.ops.quantized_decomposed.dequantize_per_tensor.default |
| 69 | ) |
| 70 | replacement_op_quant = ( |
| 71 | torch.ops.quantized_decomposed.quantize_per_tensor.default |
| 72 | ) |
| 73 | else: |
| 74 | raise ValueError(f"Invalid quantize op: {quantize.target}") |
| 75 | |
| 76 | # If user specified qparams are different from args of quantize op, we do requantization instead of eliminating quantize op |
| 77 | need_requant = False |
| 78 | if qparams is not None: |
| 79 | assert all( |
| 80 | qparam in qparams for qparam in ["scale", "zp", "dtype"] |
| 81 | ), "dtype/scale/zp must be specified in qparam for input requantization" |
| 82 | if qparams["dtype"] != quantize.args[5]: |
| 83 | if any( |
| 84 | dtype |
| 85 | not in [torch.int8, torch.uint8, torch.bool, torch.int16, torch.uint16] |
| 86 | for dtype in [qparams["dtype"], quantize.args[5]] |
| 87 | ): |