Skip to content

Show how to use QAT in TF1.x to train a network for mnist and convert it into TensorRT engine.

Notifications You must be signed in to change notification settings

shiyongming/QAT_demo

Repository files navigation

QAT_demo

Show how to use QAT in TF1.x to train a network for mnist and convert it into TensorRT engine.

Do it step by step.

  • 0.1 docker pull nvcr.io/nvidia/tensorflow:20.01-tf1-py3

  • 0.2 docker run --rm -it --gpus all nvcr.io/nvidia/tensorflow:20.01-tf1-py3

  • 0.3 git clone https://github.com/shiyongming/QAT_demo.git

  • 0.4 cd QAT_demo/

  • 0.5 pip install -r requirements.txt

  • 0.6 cd onnx-graphsurgeon && make install

  • 0.7 cd ..

  • 1 python qat_training.py

  • 2 python export_freezn_graph.py

  • 3 python fold_constants.py -i saved_results/frozen_graph.pb

  • 4

  python3 -m tf2onnx.convert --input saved_results/folded_mnist.pb\
                              --output saved_results/mnist_qat.onnx \
                              --inputs input_0:0 \
                              --outputs softmax_1:0 \
                              --opset 11 
  • 5 python postprocess_onnx.py --input saved_results/mnist_qat.onnx --output saved_results/mnist_qat_post.onnx
  • 6 python build_engine.py --onnx saved_results/mnist_qat_post.onnx --engine saved_results/mnist_qat.engine -v
  • 7 trtexec --loadEngine=saved_results/mnist_qat.engine
  • 8 python infer_checkAcc.py -e saved_results/mnist_qat.engine -b 1

The result of Step 8 looks abnormal. It may be casued by the wrong scale of input placeholder.

About

Show how to use QAT in TF1.x to train a network for mnist and convert it into TensorRT engine.

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published