quantization_speedup.rst

.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "tutorials/quantization_speedup.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        Click :ref:`here <sphx_glr_download_tutorials_quantization_speedup.py>`
        to download the full example code

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_tutorials_quantization_speedup.py:


SpeedUp Model with Calibration Config
======================================


Introduction
------------

Deep learning network has been computational intensive and memory intensive 
which increases the difficulty of deploying deep neural network model. Quantization is a 
fundamental technology which is widely used to reduce memory footprint and speedup inference 
process. Many frameworks begin to support quantization, but few of them support mixed precision 
quantization and get real speedup. Frameworks like `HAQ: Hardware-Aware Automated Quantization with Mixed Precision <https://arxiv.org/pdf/1811.08886.pdf>`__\, only support simulated mixed precision quantization which will 
not speedup the inference process. To get real speedup of mixed precision quantization and 
help people get the real feedback from hardware, we design a general framework with simple interface to allow NNI quantization algorithms to connect different 
DL model optimization backends (e.g., TensorRT, NNFusion), which gives users an end-to-end experience that after quantizing their model 
with quantization algorithms, the quantized model can be directly speeded up with the connected optimization backend. NNI connects 
TensorRT at this stage, and will support more backends in the future.


Design and Implementation
-------------------------

To support speeding up mixed precision quantization, we divide framework into two part, frontend and backend.  
Frontend could be popular training frameworks such as PyTorch, TensorFlow etc. Backend could be inference 
framework for different hardwares, such as TensorRT. At present, we support PyTorch as frontend and 
TensorRT as backend. To convert PyTorch model to TensorRT engine, we leverage onnx as intermediate graph 
representation. In this way, we convert PyTorch model to onnx model, then TensorRT parse onnx 
model to generate inference engine. 


Quantization aware training combines NNI quantization algorithm 'QAT' and NNI quantization speedup tool.
Users should set config to train quantized model using QAT algorithm(please refer to :doc:`NNI Quantization Algorithms <../compression/quantizer>`  ).
After quantization aware training, users can get new config with calibration parameters and model with quantized weight. By passing new config and model to quantization speedup tool, users can get real mixed precision speedup engine to do inference.


After getting mixed precision engine, users can do inference with input data.


Note


* Recommend using "cpu"(host) as data device(for both inference data and calibration data) since data should be on host initially and it will be transposed to device before inference. If data type is not "cpu"(host), this tool will transpose it to "cpu" which may increases unnecessary overhead.
* User can also do post-training quantization leveraging TensorRT directly(need to provide calibration dataset).
* Not all op types are supported right now. At present, NNI supports Conv, Linear, Relu and MaxPool. More op types will be supported in the following release.


Prerequisite
------------
CUDA version >= 11.0

TensorRT version >= 7.2

Note

* If you haven't installed TensorRT before or use the old version, please refer to `TensorRT Installation Guide <https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html>`__\  

Usage
-----

.. GENERATED FROM PYTHON SOURCE LINES 64-92

.. code-block:: default

    import torch
    import torch.nn.functional as F
    from torch.optim import SGD
    from scripts.compression_mnist_model import TorchModel, device, trainer, evaluator, test_trt

    config_list = [{
        'quant_types': ['input', 'weight'],
        'quant_bits': {'input': 8, 'weight': 8},
        'op_types': ['Conv2d']
    }, {
        'quant_types': ['output'],
        'quant_bits': {'output': 8},
        'op_types': ['ReLU']
    }, {
        'quant_types': ['input', 'weight'],
        'quant_bits': {'input': 8, 'weight': 8},
        'op_names': ['fc1', 'fc2']
    }]

    model = TorchModel().to(device)
    optimizer = SGD(model.parameters(), lr=0.01, momentum=0.5)
    criterion = F.nll_loss
    dummy_input = torch.rand(32, 1, 28, 28).to(device)

    from nni.algorithms.compression.pytorch.quantization import QAT_Quantizer
    quantizer = QAT_Quantizer(model, config_list, optimizer, dummy_input)
    quantizer.compress()





.. rst-class:: sphx-glr-script-out

 Out:

 .. code-block:: none


    TorchModel(
      (conv1): QuantizerModuleWrapper(
        (module): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1))
      )
      (conv2): QuantizerModuleWrapper(
        (module): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
      )
      (fc1): QuantizerModuleWrapper(
        (module): Linear(in_features=256, out_features=120, bias=True)
      )
      (fc2): QuantizerModuleWrapper(
        (module): Linear(in_features=120, out_features=84, bias=True)
      )
      (fc3): Linear(in_features=84, out_features=10, bias=True)
      (relu1): QuantizerModuleWrapper(
        (module): ReLU()
      )
      (relu2): QuantizerModuleWrapper(
        (module): ReLU()
      )
      (relu3): QuantizerModuleWrapper(
        (module): ReLU()
      )
      (relu4): QuantizerModuleWrapper(
        (module): ReLU()
      )
      (pool1): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), padding=0, dilation=1, ceil_mode=False)
      (pool2): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), padding=0, dilation=1, ceil_mode=False)
    )



.. GENERATED FROM PYTHON SOURCE LINES 93-94

finetuning the model by using QAT

.. GENERATED FROM PYTHON SOURCE LINES 94-98

.. code-block:: default

    for epoch in range(3):
        trainer(model, optimizer, criterion)
        evaluator(model)





.. rst-class:: sphx-glr-script-out

 Out:

 .. code-block:: none

    Average test loss: 0.5386, Accuracy: 8619/10000 (86%)
    Average test loss: 0.1553, Accuracy: 9521/10000 (95%)
    Average test loss: 0.1001, Accuracy: 9686/10000 (97%)




.. GENERATED FROM PYTHON SOURCE LINES 99-100

export model and get calibration_config

.. GENERATED FROM PYTHON SOURCE LINES 100-108

.. code-block:: default

    import os
    os.makedirs('log', exist_ok=True)
    model_path = "./log/mnist_model.pth"
    calibration_path = "./log/mnist_calibration.pth"
    calibration_config = quantizer.export_model(model_path, calibration_path)

    print("calibration_config: ", calibration_config)





.. rst-class:: sphx-glr-script-out

 Out:

 .. code-block:: none

    calibration_config:  {'conv1': {'weight_bits': 8, 'weight_scale': tensor([0.0029], device='cuda:0'), 'weight_zero_point': tensor([98.], device='cuda:0'), 'input_bits': 8, 'tracked_min_input': -0.4242129623889923, 'tracked_max_input': 2.821486711502075}, 'conv2': {'weight_bits': 8, 'weight_scale': tensor([0.0017], device='cuda:0'), 'weight_zero_point': tensor([124.], device='cuda:0'), 'input_bits': 8, 'tracked_min_input': 0.0, 'tracked_max_input': 8.848002433776855}, 'fc1': {'weight_bits': 8, 'weight_scale': tensor([0.0010], device='cuda:0'), 'weight_zero_point': tensor([134.], device='cuda:0'), 'input_bits': 8, 'tracked_min_input': 0.0, 'tracked_max_input': 14.64758586883545}, 'fc2': {'weight_bits': 8, 'weight_scale': tensor([0.0013], device='cuda:0'), 'weight_zero_point': tensor([121.], device='cuda:0'), 'input_bits': 8, 'tracked_min_input': 0.0, 'tracked_max_input': 15.807988166809082}, 'relu1': {'output_bits': 8, 'tracked_min_output': 0.0, 'tracked_max_output': 9.041301727294922}, 'relu2': {'output_bits': 8, 'tracked_min_output': 0.0, 'tracked_max_output': 15.143928527832031}, 'relu3': {'output_bits': 8, 'tracked_min_output': 0.0, 'tracked_max_output': 16.151935577392578}, 'relu4': {'output_bits': 8, 'tracked_min_output': 0.0, 'tracked_max_output': 11.749024391174316}}




.. GENERATED FROM PYTHON SOURCE LINES 109-110

build tensorRT engine to make a real speedup

.. GENERATED FROM PYTHON SOURCE LINES 110-117

.. code-block:: default


    from nni.compression.pytorch.quantization_speedup import ModelSpeedupTensorRT
    input_shape = (32, 1, 28, 28)
    engine = ModelSpeedupTensorRT(model, input_shape, config=calibration_config, batchsize=32)
    engine.compress()
    test_trt(engine)





.. rst-class:: sphx-glr-script-out

 Out:

 .. code-block:: none

    Loss: 0.10061546401977539  Accuracy: 96.83%
    Inference elapsed_time (whole dataset): 0.04322671890258789s




.. GENERATED FROM PYTHON SOURCE LINES 118-169

Note that NNI also supports post-training quantization directly, please refer to complete examples for detail.

For complete examples please refer to :githublink:`the code <examples/model_compress/quantization/mixed_precision_speedup_mnist.py>`.

For more parameters about the class 'TensorRTModelSpeedUp', you can refer to :doc:`Model Compression API Reference <../reference/compression/quantization_speedup>`.

Mnist test
^^^^^^^^^^

on one GTX2080 GPU,
input tensor: ``torch.randn(128, 1, 28, 28)``

.. list-table::
   :header-rows: 1
   :widths: auto

   * - quantization strategy
     - Latency
     - accuracy
   * - all in 32bit
     - 0.001199961
     - 96%
   * - mixed precision(average bit 20.4)
     - 0.000753688
     - 96%
   * - all in 8bit
     - 0.000229869
     - 93.7%

Cifar10 resnet18 test (train one epoch)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

on one GTX2080 GPU,
input tensor: ``torch.randn(128, 3, 32, 32)``

.. list-table::
   :header-rows: 1
   :widths: auto

   * - quantization strategy
     - Latency
     - accuracy
   * - all in 32bit
     - 0.003286268
     - 54.21%
   * - mixed precision(average bit 11.55)
     - 0.001358022
     - 54.78%
   * - all in 8bit
     - 0.000859139
     - 52.81%


.. rst-class:: sphx-glr-timing

   **Total running time of the script:** ( 1 minutes  4.509 seconds)


.. _sphx_glr_download_tutorials_quantization_speedup.py:


.. only :: html

 .. container:: sphx-glr-footer
    :class: sphx-glr-footer-example



  .. container:: sphx-glr-download sphx-glr-download-python

     :download:`Download Python source code: quantization_speedup.py <quantization_speedup.py>`



  .. container:: sphx-glr-download sphx-glr-download-jupyter

     :download:`Download Jupyter notebook: quantization_speedup.ipynb <quantization_speedup.ipynb>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_