Returns quantization parameters such as scale/zero_point: { "inputs": { : {"scale": float, "zero_point": int} }, "outputs": { : {"scale": float, "zero_point": int} } } Note that this function will
(
edge_prog: EdgeProgramManager,
*,
input_idxs: Sequence[int] = (0,),
output_idxs: Sequence[int] = (0,),
)
| 325 | |
| 326 | |
| 327 | def extract_io_quant_params( |
| 328 | edge_prog: EdgeProgramManager, |
| 329 | *, |
| 330 | input_idxs: Sequence[int] = (0,), |
| 331 | output_idxs: Sequence[int] = (0,), |
| 332 | ) -> Dict[str, Dict[str, Dict[str, Any]]]: |
| 333 | """ |
| 334 | Returns quantization parameters such as scale/zero_point: |
| 335 | { |
| 336 | "inputs": { |
| 337 | <placeholder_name>: {"scale": float, "zero_point": int} |
| 338 | }, |
| 339 | "outputs": { |
| 340 | <node_name>: {"scale": float, "zero_point": int} |
| 341 | } |
| 342 | } |
| 343 | |
| 344 | Note that this function will strip out the IO quantize/dequantize ops as |
| 345 | it records their parameters, so if you need to preserve the original graph |
| 346 | you need to make a copy with copy.deepcopy before. |
| 347 | |
| 348 | Note that `to_edge_transform_and_lower` should be called before. |
| 349 | """ |
| 350 | # Use IO passes |
| 351 | passes = [] |
| 352 | for idx in input_idxs: |
| 353 | passes.append(QuantizeInputs(edge_prog, [idx])) |
| 354 | for idx in output_idxs: |
| 355 | passes.append(QuantizeOutputs(edge_prog, [idx])) |
| 356 | |
| 357 | # Apply them |
| 358 | edge_prog = edge_prog.transform(passes) |
| 359 | |
| 360 | cfg = getattr(edge_prog, "_config_methods", {}) or {} |
| 361 | |
| 362 | # We need GraphModule to find node names |
| 363 | gm = edge_prog.exported_program().graph_module |
| 364 | |
| 365 | input_names = _gather_io_names(gm, side="input") |
| 366 | output_names = _gather_io_names(gm, side="output") |
| 367 | |
| 368 | # Build the result dict |
| 369 | result = {"inputs": {}, "outputs": {}} |
| 370 | for key, val in cfg.items(): |
| 371 | if key.startswith("input"): |
| 372 | prefix, section, names = "input", "inputs", input_names |
| 373 | elif key.startswith("output"): |
| 374 | prefix, section, names = "output", "outputs", output_names |
| 375 | else: |
| 376 | continue |
| 377 | |
| 378 | idx_str, param = key[len(prefix) :].split("_", 1) |
| 379 | idx = int(idx_str) |
| 380 | name = names[idx] |
| 381 | # We need to map 'zp' to 'zero_point' |
| 382 | out_param = "zero_point" if param in ("zp", "zero_point") else param |
| 383 | result[section].setdefault(name, {})[out_param] = val |
| 384 |