次のコードは、vai_q_tensorflow2 API を使用してトレーニング後の量子化を実行し、量子化されたモデルを ONNX にエクスポートする方法を示しています。
model = tf.keras.models.load_model('float_model.h5')
from tensorflow_model_optimization.quantization.keras import vitis_quantize
quantizer = vitis_quantize.VitisQuantizer(model)
quantized_model = quantizer.quantize_model(calib_dataset=calib_dataset,
output_format='onnx',
onnx_opset_version=11,
output_dir='./quantize_results',
**kwargs)
- output_format
-
文字列。量子化されたモデルをどの形式で保存するかを指定します。オプションは次のとおりです。
- '': 保存をスキップ
- h5: .h5 ファイルを保存
- tf: saved_model ファイルを保存
- onnx: ONNX ファイルを保存
デフォルト値は「''」です。
- onnx_opset_version
- Int。ONNX opset バージョンを指定します。output_format が 'onnx' の場合のみ有効です。デフォルト値は 11 です。
- output_dir
- 文字列。量子化されたモデルを保存するディレクトリを指定します。デフォルト値は「./quantize_results」です。