| 27 | |
| 28 | |
| 29 | def _onnx( |
| 30 | model, |
| 31 | quantize: bool = False, |
| 32 | opset_version: int = 14, |
| 33 | export_dir: str = None, |
| 34 | **kwargs, |
| 35 | ): |
| 36 | |
| 37 | dummy_input = model.export_dummy_inputs() |
| 38 | |
| 39 | verbose = kwargs.get("verbose", False) |
| 40 | |
| 41 | export_name = model.export_name() |
| 42 | model_path = os.path.join(export_dir, export_name) |
| 43 | torch.onnx.export( |
| 44 | model, |
| 45 | dummy_input, |
| 46 | model_path, |
| 47 | verbose=verbose, |
| 48 | opset_version=opset_version, |
| 49 | input_names=model.export_input_names(), |
| 50 | output_names=model.export_output_names(), |
| 51 | dynamic_axes=model.export_dynamic_axes(), |
| 52 | ) |
| 53 | |
| 54 | if quantize: |
| 55 | from onnxruntime.quantization import QuantType, quantize_dynamic |
| 56 | import onnx |
| 57 | |
| 58 | quant_model_path = model_path.replace(".onnx", "_quant.onnx") |
| 59 | if not os.path.exists(quant_model_path): |
| 60 | onnx_model = onnx.load(model_path) |
| 61 | nodes = [n.name for n in onnx_model.graph.node] |
| 62 | nodes_to_exclude = [ |
| 63 | m for m in nodes if "output" in m or "bias_encoder" in m or "bias_decoder" in m |
| 64 | ] |
| 65 | quantize_dynamic( |
| 66 | model_input=model_path, |
| 67 | model_output=quant_model_path, |
| 68 | op_types_to_quantize=["MatMul"], |
| 69 | per_channel=True, |
| 70 | reduce_range=False, |
| 71 | weight_type=QuantType.QUInt8, |
| 72 | nodes_to_exclude=nodes_to_exclude, |
| 73 | ) |