| # import onnx | |
| # # Load the ONNX model | |
| # model_path = "model_uint8.onnx" # Replace with the path to your ONNX model | |
| # onnx_model = onnx.load(model_path) | |
| # # Print model's input and output shapes | |
| # for input_tensor in onnx_model.graph.input: | |
| # print(f"Input Name: {input_tensor.name}") | |
| # print( | |
| # f"Input Shape: {[dim.dim_value for dim in input_tensor.type.tensor_type.shape.dim]}" | |
| # ) | |
| # for output_tensor in onnx_model.graph.output: | |
| # print(f"Output Name: {output_tensor.name}") | |
| # print( | |
| # f"Output Shape: {[dim.dim_value for dim in output_tensor.type.tensor_type.shape.dim]}" | |
| # ) | |
| from onnxruntime.quantization import quantize_dynamic, QuantType | |
| # Define the path to the original ONNX model and the quantized output model | |
| onnx_model_path = "./granite_embedding_model.onnx" # Path to the original ONNX model | |
| quantized_model_path = "./model_uint8.onnx" # Path to save the quantized ONNX model | |
| # Perform dynamic quantization to UInt8 | |
| quantize_dynamic( | |
| model_input=onnx_model_path, # Input ONNX model path | |
| model_output=quantized_model_path, # Output quantized model path | |
| weight_type=QuantType.QUInt8, # Use UInt8 for weights | |
| ) | |
| # Print confirmation of quantization | |
| print(f"Quantized model saved to {quantized_model_path}") | |