The example/lstm_quant_tensorflow/quantize_lstm.py file contains an
example.
- Import the TensorFlow quantizer
module.
from tf_nndct.quantization.api import tf_quantizer
- Generate a quantizer with quantization needed input, and the batch size of
input data must be 1. Get the converted
model.
single_batch_data = X_test[:1, ] input_signature = tf.TensorSpec(single_batch_data.shape[1:], tf.int32) quantizer = tf_quantizer(model, input_signature, quant_mode=args.quant_mode, bitwidth=16) rebuilt_model = quantizer.quant_model
- Forward a neural network with the converted
model.
output = rebuilt_model(X_test[start: end])
- Output the quantization result and deploy the model. When dumping the outputs,
the batch size of the input data must be 1.
if args.quant_mode == 'calib': quantizer.export_quant_config() elif args.quant_mode == 'test': quantizer.dump_xmodel() quantizer.dump_rnn_outputs_by_timestep(X_test[:1])