Implementation quantization-aware training on classification model with Tensorflow and Keras.
(Keras was be included by Tensorflow in this version)
tensorflow - 1.15.0
Ubuntu - 16.04
It is based on Tensorflow 1.15.0, it can help your CNN model to do quantization-aware training simply, all you need to do is prepare your Keras model and dataset.
The quantization-aware training will transform int8 from float32.Int8 can be compiled on Edge TPU or mobile device. It uses fake quantization nodes to simulate the effect of quantization when you're training. The code will use quantization-aware training to produce checkpoints(.data, .mete, .index) in checkpoint folder, and then export them into the new .pd file.
-
You can refer to model/example_VGG/VGG.py to choose the learning rate , optimizer to create your own model (save as model.h5).
-
Prepare your data in data folder:
train--
-class1
-class2
- 1.jpg
- 2.jpg
'
'
'
test--
-class1
-class2
- 1.jpg
'
'
'
val--
-class1
'
'
'
- Try the quantize model works by running it through the main.py script:
python main.py \
--model_path your/model/path \
--epoch ChooseYourEpoch \
--batch_size ChooseYourBatchSize
- You can get the .pd file that your model has been quantized
python main.py \
--model_path ./model/example_VGG/model.h5 \
--epoch 10 \
--batch_size 16