[Doc] Compression (#4574)

docs/source/tutorials.rst

@@ -10,6 +10,10 @@ Tutorials
    tutorials/nni_experiment
    tutorials/hello_nas
    tutorials/nasbench_as_dataset
+   tutorials/pruning_quick_start_mnist
+   tutorials/pruning_speed_up
+   tutorials/quantization_quick_start_mnist
+   tutorials/quantization_speed_up
 
 .. ----------------------
 
@@ -35,3 +39,35 @@ Tutorials
    :image: ../img/thumbnails/overview-30.png
    :background: pink
    :tags: NAS
+
+.. cardlinkitem::
+   :header: Get Started with Model Pruning on MNIST
+   :description: Familiarize yourself with pruning to compress your model 
+   :link: tutorials/pruning_quick_start_mnist.html
+   :image: ../img/thumbnails/overview-29.png
+   :background: cyan
+   :tags: Compression
+
+.. cardlinkitem::
+   :header: Get Started with Model Quantization on MNIST
+   :description: Familiarize yourself with quantization to compress your model
+   :link: tutorials/quantization_quick_start_mnist.html
+   :image: ../img/thumbnails/overview-29.png
+   :background: cyan
+   :tags: Compression
+
+.. cardlinkitem::
+   :header: Speed Up Model with Mask
+   :description: Make your model real smaller and faster with speed-up after pruned by pruner
+   :link: tutorials/pruning_speed_up.html
+   :image: ../img/thumbnails/overview-29.png
+   :background: cyan
+   :tags: Compression
+
+.. cardlinkitem::
+   :header: Speed Up Model with Calibration Config
+   :description: Make your model real smaller and faster with speed-up after quantized by quantizer
+   :link: tutorials/quantization_speed_up.html
+   :image: ../img/thumbnails/overview-29.png
+   :background: cyan
+   :tags: Compression

docs/source/tutorials/index.rst

@@ -9,6 +9,27 @@ Tutorials
 
 
 
+.. raw:: html
+
+    <div class="sphx-glr-thumbcontainer" tooltip="Introduction ------------">
+
+.. only:: html
+
+ .. figure:: /tutorials/images/thumb/sphx_glr_pruning_speed_up_thumb.png
+     :alt: Speed Up Model with Mask
+
+     :ref:`sphx_glr_tutorials_pruning_speed_up.py`
+
+.. raw:: html
+
+    </div>
+
+
+.. toctree::
+   :hidden:
+
+   /tutorials/pruning_speed_up
+
 .. raw:: html
 
     <div class="sphx-glr-thumbcontainer" tooltip="Start and Manage a New Experiment">
@@ -30,6 +51,69 @@ Tutorials
 
    /tutorials/nni_experiment
 
+.. raw:: html
+
+    <div class="sphx-glr-thumbcontainer" tooltip="Model pruning is a technique to reduce the model size and computation by reducing model weight ...">
+
+.. only:: html
+
+ .. figure:: /tutorials/images/thumb/sphx_glr_pruning_quick_start_mnist_thumb.png
+     :alt: Pruning Quickstart
+
+     :ref:`sphx_glr_tutorials_pruning_quick_start_mnist.py`
+
+.. raw:: html
+
+    </div>
+
+
+.. toctree::
+   :hidden:
+
+   /tutorials/pruning_quick_start_mnist
+
+.. raw:: html
+
+    <div class="sphx-glr-thumbcontainer" tooltip="Quantization reduces model size and speeds up inference time by reducing the number of bits req...">
+
+.. only:: html
+
+ .. figure:: /tutorials/images/thumb/sphx_glr_quantization_quick_start_mnist_thumb.png
+     :alt: Quantization Quickstart
+
+     :ref:`sphx_glr_tutorials_quantization_quick_start_mnist.py`
+
+.. raw:: html
+
+    </div>
+
+
+.. toctree::
+   :hidden:
+
+   /tutorials/quantization_quick_start_mnist
+
+.. raw:: html
+
+    <div class="sphx-glr-thumbcontainer" tooltip=" Introduction ------------">
+
+.. only:: html
+
+ .. figure:: /tutorials/images/thumb/sphx_glr_quantization_speed_up_thumb.png
+     :alt: Speed Up Model with Calibration Config
+
+     :ref:`sphx_glr_tutorials_quantization_speed_up.py`
+
+.. raw:: html
+
+    </div>
+
+
+.. toctree::
+   :hidden:
+
+   /tutorials/quantization_speed_up
+
 .. raw:: html
 
     <div class="sphx-glr-thumbcontainer" tooltip="In this tutorial, we show how to use NAS Benchmarks as datasets. For research purposes we somet...">

docs/source/tutorials/pruning_quick_start_mnist.ipynb

+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "%matplotlib inline"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "\n# Pruning Quickstart\n\nModel pruning is a technique to reduce the model size and computation by reducing model weight size or intermediate state size.\nIt usually has following paths:\n\n#. Pre-training a model -> Pruning the model -> Fine-tuning the model\n#. Pruning the model aware training -> Fine-tuning the model\n#. Pruning the model -> Pre-training the compact model\n\nNNI supports the above three modes and mainly focuses on the pruning stage.\nFollow this tutorial for a quick look at how to use NNI to prune a model in a common practice.\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Preparation\n\nIn this tutorial, we use a simple model and pre-train on MNIST dataset.\nIf you are familiar with defining a model and training in pytorch, you can skip directly to `Pruning Model`_.\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "import torch\nimport torch.nn.functional as F\nfrom torch.optim import SGD\n\nfrom scripts.compression_mnist_model import TorchModel, trainer, evaluator, device\n\n# define the model\nmodel = TorchModel().to(device)\n\n# define the optimizer and criterion for pre-training\n\noptimizer = SGD(model.parameters(), 1e-2)\ncriterion = F.nll_loss\n\n# pre-train and evaluate the model on MNIST dataset\nfor epoch in range(3):\n    trainer(model, optimizer, criterion)\n    evaluator(model)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Pruning Model\n\nUsing L1NormPruner pruning the model and generating the masks.\nUsually, pruners require original model and ``config_list`` as parameters.\nDetailed about how to write ``config_list`` please refer ...\n\nThis `config_list` means all layers whose type is `Linear` or `Conv2d` will be pruned,\nexcept the layer named `fc3`, because `fc3` is `exclude`.\nThe final sparsity ratio for each layer is 50%. The layer named `fc3` will not be pruned.\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "config_list = [{\n    'sparsity_per_layer': 0.5,\n    'op_types': ['Linear', 'Conv2d']\n}, {\n    'exclude': True,\n    'op_names': ['fc3']\n}]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "Pruners usually require `model` and `config_list` as input arguments.\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "from nni.algorithms.compression.v2.pytorch.pruning import L1NormPruner\n\npruner = L1NormPruner(model, config_list)\n# show the wrapped model structure\nprint(model)\n# compress the model and generate the masks\n_, masks = pruner.compress()\n# show the masks sparsity\nfor name, mask in masks.items():\n    print(name, ' sparsity: ', '{:.2}'.format(mask['weight'].sum() / mask['weight'].numel()))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "Speed up the original model with masks, note that `ModelSpeedup` requires an unwrapped model.\nThe model becomes smaller after speed-up,\nand reaches a higher sparsity ratio because `ModelSpeedup` will propagate the masks across layers.\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "# need to unwrap the model, if the model is wrapped before speed up\npruner._unwrap_model()\n\n# speed up the model\nfrom nni.compression.pytorch.speedup import ModelSpeedup\n\nModelSpeedup(model, torch.rand(3, 1, 28, 28).to(device), masks).speedup_model()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "the model will become real smaller after speed up\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "print(model)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Fine-tuning Compacted Model\nNote that if the model has been sped up, you need to re-initialize a new optimizer for fine-tuning.\nBecause speed up will replace the masked big layers with dense small ones.\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "optimizer = SGD(model.parameters(), 1e-2)\nfor epoch in range(3):\n    trainer(model, optimizer, criterion)"
+      ]
+    }
+  ],
+  "metadata": {
+    "kernelspec": {
+      "display_name": "Python 3",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "codemirror_mode": {
+        "name": "ipython",
+        "version": 3
+      },
+      "file_extension": ".py",
+      "mimetype": "text/x-python",
+      "name": "python",
+      "nbconvert_exporter": "python",
+      "pygments_lexer": "ipython3",
+      "version": "3.8.8"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
\ No newline at end of file

docs/source/tutorials/pruning_quick_start_mnist.py

+"""
+Pruning Quickstart
+==================
+
+Model pruning is a technique to reduce the model size and computation by reducing model weight size or intermediate state size.
+It usually has following paths:
+
+#. Pre-training a model -> Pruning the model -> Fine-tuning the model
+#. Pruning the model aware training -> Fine-tuning the model
+#. Pruning the model -> Pre-training the compact model
+
+NNI supports the above three modes and mainly focuses on the pruning stage.
+Follow this tutorial for a quick look at how to use NNI to prune a model in a common practice.
+"""
+
+# %%
+# Preparation
+# -----------
+#
+# In this tutorial, we use a simple model and pre-train on MNIST dataset.
+# If you are familiar with defining a model and training in pytorch, you can skip directly to `Pruning Model`_.
+
+import torch
+import torch.nn.functional as F
+from torch.optim import SGD
+
+from scripts.compression_mnist_model import TorchModel, trainer, evaluator, device
+
+# define the model
+model = TorchModel().to(device)
+
+# define the optimizer and criterion for pre-training
+
+optimizer = SGD(model.parameters(), 1e-2)
+criterion = F.nll_loss
+
+# pre-train and evaluate the model on MNIST dataset
+for epoch in range(3):
+    trainer(model, optimizer, criterion)
+    evaluator(model)
+
+# %%
+# Pruning Model
+# -------------
+#
+# Using L1NormPruner pruning the model and generating the masks.
+# Usually, pruners require original model and ``config_list`` as parameters.
+# Detailed about how to write ``config_list`` please refer ...
+#
+# This `config_list` means all layers whose type is `Linear` or `Conv2d` will be pruned,
+# except the layer named `fc3`, because `fc3` is `exclude`.
+# The final sparsity ratio for each layer is 50%. The layer named `fc3` will not be pruned.
+
+config_list = [{
+    'sparsity_per_layer': 0.5,
+    'op_types': ['Linear', 'Conv2d']
+}, {
+    'exclude': True,
+    'op_names': ['fc3']
+}]
+
+# %%
+# Pruners usually require `model` and `config_list` as input arguments.
+
+from nni.algorithms.compression.v2.pytorch.pruning import L1NormPruner
+
+pruner = L1NormPruner(model, config_list)
+# show the wrapped model structure
+print(model)
+# compress the model and generate the masks
+_, masks = pruner.compress()
+# show the masks sparsity
+for name, mask in masks.items():
+    print(name, ' sparsity: ', '{:.2}'.format(mask['weight'].sum() / mask['weight'].numel()))
+
+# %%
+# Speed up the original model with masks, note that `ModelSpeedup` requires an unwrapped model.
+# The model becomes smaller after speed-up,
+# and reaches a higher sparsity ratio because `ModelSpeedup` will propagate the masks across layers.
+
+# need to unwrap the model, if the model is wrapped before speed up
+pruner._unwrap_model()
+
+# speed up the model
+from nni.compression.pytorch.speedup import ModelSpeedup
+
+ModelSpeedup(model, torch.rand(3, 1, 28, 28).to(device), masks).speedup_model()
+
+# %%
+# the model will become real smaller after speed up
+print(model)
+
+# %%
+# Fine-tuning Compacted Model
+# ---------------------------
+# Note that if the model has been sped up, you need to re-initialize a new optimizer for fine-tuning.
+# Because speed up will replace the masked big layers with dense small ones.
+
+optimizer = SGD(model.parameters(), 1e-2)
+for epoch in range(3):
+    trainer(model, optimizer, criterion)

docs/source/tutorials/pruning_quick_start_mnist.py.md5

+775624e7a28ae5c6eb2027eace7fff67
\ No newline at end of file

docs/source/tutorials/pruning_quick_start_mnist.rst

+
+.. DO NOT EDIT.
+.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
+.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
+.. "tutorials/pruning_quick_start_mnist.py"
+.. LINE NUMBERS ARE GIVEN BELOW.
+
+.. only:: html
+
+    .. note::
+        :class: sphx-glr-download-link-note
+
+        Click :ref:`here <sphx_glr_download_tutorials_pruning_quick_start_mnist.py>`
+        to download the full example code
+
+.. rst-class:: sphx-glr-example-title
+
+.. _sphx_glr_tutorials_pruning_quick_start_mnist.py:
+
+
+Pruning Quickstart
+==================
+
+Model pruning is a technique to reduce the model size and computation by reducing model weight size or intermediate state size.
+It usually has following paths:
+
+#. Pre-training a model -> Pruning the model -> Fine-tuning the model
+#. Pruning the model aware training -> Fine-tuning the model
+#. Pruning the model -> Pre-training the compact model
+
+NNI supports the above three modes and mainly focuses on the pruning stage.
+Follow this tutorial for a quick look at how to use NNI to prune a model in a common practice.
+
+.. GENERATED FROM PYTHON SOURCE LINES 17-22
+
+Preparation
+-----------
+
+In this tutorial, we use a simple model and pre-train on MNIST dataset.
+If you are familiar with defining a model and training in pytorch, you can skip directly to `Pruning Model`_.
+
+.. GENERATED FROM PYTHON SOURCE LINES 22-42
+
+.. code-block:: default
+
+
+    import torch
+    import torch.nn.functional as F
+    from torch.optim import SGD
+
+    from scripts.compression_mnist_model import TorchModel, trainer, evaluator, device
+
+    # define the model
+    model = TorchModel().to(device)
+
+    # define the optimizer and criterion for pre-training
+
+    optimizer = SGD(model.parameters(), 1e-2)
+    criterion = F.nll_loss
+
+    # pre-train and evaluate the model on MNIST dataset
+    for epoch in range(3):
+        trainer(model, optimizer, criterion)
+        evaluator(model)
+
+
+
+
+
+.. rst-class:: sphx-glr-script-out
+
+ Out:
+
+ .. code-block:: none
+
+    Average test loss: 1.8381, Accuracy: 5939/10000 (59%)
+    Average test loss: 0.3143, Accuracy: 9045/10000 (90%)
+    Average test loss: 0.1928, Accuracy: 9387/10000 (94%)
+
+
+
+
+.. GENERATED FROM PYTHON SOURCE LINES 43-53
+
+Pruning Model
+-------------
+
+Using L1NormPruner pruning the model and generating the masks.
+Usually, pruners require original model and ``config_list`` as parameters.
+Detailed about how to write ``config_list`` please refer ...
+
+This `config_list` means all layers whose type is `Linear` or `Conv2d` will be pruned,
+except the layer named `fc3`, because `fc3` is `exclude`.
+The final sparsity ratio for each layer is 50%. The layer named `fc3` will not be pruned.
+
+.. GENERATED FROM PYTHON SOURCE LINES 53-62
+
+.. code-block:: default
+
+
+    config_list = [{
+        'sparsity_per_layer': 0.5,
+        'op_types': ['Linear', 'Conv2d']
+    }, {
+        'exclude': True,
+        'op_names': ['fc3']
+    }]
+
+
+
+
+
+
+
+
+.. GENERATED FROM PYTHON SOURCE LINES 63-64
+
+Pruners usually require `model` and `config_list` as input arguments.
+
+.. GENERATED FROM PYTHON SOURCE LINES 64-76
+
+.. code-block:: default
+
+
+    from nni.algorithms.compression.v2.pytorch.pruning import L1NormPruner
+
+    pruner = L1NormPruner(model, config_list)
+    # show the wrapped model structure
+    print(model)
+    # compress the model and generate the masks
+    _, masks = pruner.compress()
+    # show the masks sparsity
+    for name, mask in masks.items():
+        print(name, ' sparsity: ', '{:.2}'.format(mask['weight'].sum() / mask['weight'].numel()))
+
+
+
+
+
+.. rst-class:: sphx-glr-script-out
+
+ Out:
+
+ .. code-block:: none
+
+    TorchModel(
+      (conv1): PrunerModuleWrapper(
+        (module): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1))
+      )
+      (conv2): PrunerModuleWrapper(
+        (module): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
+      )
+      (fc1): PrunerModuleWrapper(
+        (module): Linear(in_features=256, out_features=120, bias=True)
+      )
+      (fc2): PrunerModuleWrapper(
+        (module): Linear(in_features=120, out_features=84, bias=True)
+      )
+      (fc3): Linear(in_features=84, out_features=10, bias=True)
+    )
+    conv1  sparsity:  0.5
+    conv2  sparsity:  0.5
+    fc1  sparsity:  0.5
+    fc2  sparsity:  0.5
+
+
+
+
+.. GENERATED FROM PYTHON SOURCE LINES 77-80
+
+Speed up the original model with masks, note that `ModelSpeedup` requires an unwrapped model.
+The model becomes smaller after speed-up,
+and reaches a higher sparsity ratio because `ModelSpeedup` will propagate the masks across layers.
+
+.. GENERATED FROM PYTHON SOURCE LINES 80-89
+
+.. code-block:: default
+
+
+    # need to unwrap the model, if the model is wrapped before speed up
+    pruner._unwrap_model()
+
+    # speed up the model
+    from nni.compression.pytorch.speedup import ModelSpeedup
+
+    ModelSpeedup(model, torch.rand(3, 1, 28, 28).to(device), masks).speedup_model()
+
+
+
+
+
+.. rst-class:: sphx-glr-script-out
+
+ Out:
+
+ .. code-block:: none
+
+    /home/ningshang/nni/nni/compression/pytorch/utils/mask_conflict.py:124: UserWarning: This overload of nonzero is deprecated:
+            nonzero()
+    Consider using one of the following signatures instead:
+            nonzero(*, bool as_tuple) (Triggered internally at  /pytorch/torch/csrc/utils/python_arg_parser.cpp:766.)
+      all_ones = (w_mask.flatten(1).sum(-1) == count).nonzero().squeeze(1).tolist()
+    /home/ningshang/nni/nni/compression/pytorch/speedup/infer_mask.py:262: UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the gradient for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more informations.
+      if isinstance(self.output, torch.Tensor) and self.output.grad is not None:
+    /home/ningshang/nni/nni/compression/pytorch/speedup/compressor.py:282: UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the gradient for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more informations.
+      if last_output.grad is not None and tin.grad is not None:
+
+
+
+
+.. GENERATED FROM PYTHON SOURCE LINES 90-91
+
+the model will become real smaller after speed up
+
+.. GENERATED FROM PYTHON SOURCE LINES 91-93
+
+.. code-block:: default
+
+    print(model)
+
+
+
+
+
+.. rst-class:: sphx-glr-script-out
+
+ Out:
+
+ .. code-block:: none
+
+    TorchModel(
+      (conv1): Conv2d(1, 3, kernel_size=(5, 5), stride=(1, 1))
+      (conv2): Conv2d(3, 8, kernel_size=(5, 5), stride=(1, 1))
+      (fc1): Linear(in_features=128, out_features=60, bias=True)
+      (fc2): Linear(in_features=60, out_features=42, bias=True)
+      (fc3): Linear(in_features=42, out_features=10, bias=True)
+    )
+
+
+
+
+.. GENERATED FROM PYTHON SOURCE LINES 94-98
+
+Fine-tuning Compacted Model
+---------------------------
+Note that if the model has been sped up, you need to re-initialize a new optimizer for fine-tuning.
+Because speed up will replace the masked big layers with dense small ones.
+
+.. GENERATED FROM PYTHON SOURCE LINES 98-102
+
+.. code-block:: default
+
+
+    optimizer = SGD(model.parameters(), 1e-2)
+    for epoch in range(3):
+        trainer(model, optimizer, criterion)
+
+
+
+
+
+
+
+
+.. rst-class:: sphx-glr-timing
+
+   **Total running time of the script:** ( 1 minutes  15.845 seconds)
+
+
+.. _sphx_glr_download_tutorials_pruning_quick_start_mnist.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: pruning_quick_start_mnist.py <pruning_quick_start_mnist.py>`
+
+
+
+  .. container:: sphx-glr-download sphx-glr-download-jupyter
+
+     :download:`Download Jupyter notebook: pruning_quick_start_mnist.ipynb <pruning_quick_start_mnist.ipynb>`
+
+
+.. only:: html
+
+ .. rst-class:: sphx-glr-signature
+
+    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_

docs/source/tutorials/pruning_speed_up.py.md5

+5bcdee7241d8daf931bd76f435167a58
\ No newline at end of file

docs/source/tutorials/quantization_quick_start_mnist.ipynb

+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "%matplotlib inline"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "\n# Quantization Quickstart\n\nQuantization reduces model size and speeds up inference time by reducing the number of bits required to represent weights or activations.\n\nIn NNI, both post-training quantization algorithms and quantization-aware training algorithms are supported.\nHere we use `QAT_Quantizer` as an example to show the usage of quantization in NNI.\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Preparation\n\nIn this tutorial, we use a simple model and pre-train on MNIST dataset.\nIf you are familiar with defining a model and training in pytorch, you can skip directly to `Quantizing Model`_.\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "import torch\nimport torch.nn.functional as F\nfrom torch.optim import SGD\n\nfrom scripts.compression_mnist_model import TorchModel, trainer, evaluator, device\n\n# define the model\nmodel = TorchModel().to(device)\n\n# define the optimizer and criterion for pre-training\n\noptimizer = SGD(model.parameters(), 1e-2)\ncriterion = F.nll_loss\n\n# pre-train and evaluate the model on MNIST dataset\nfor epoch in range(3):\n    trainer(model, optimizer, criterion)\n    evaluator(model)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Quantizing Model\n\nInitialize a `config_list`.\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "config_list = [{\n    'quant_types': ['input', 'weight'],\n    'quant_bits': {'input': 8, 'weight': 8},\n    'op_names': ['conv1']\n}, {\n    'quant_types': ['output'],\n    'quant_bits': {'output': 8},\n    'op_names': ['relu1']\n}, {\n    'quant_types': ['input', 'weight'],\n    'quant_bits': {'input': 8, 'weight': 8},\n    'op_names': ['conv2']\n}, {\n    'quant_types': ['output'],\n    'quant_bits': {'output': 8},\n    'op_names': ['relu2']\n}]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "finetuning the model by using QAT\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "from nni.algorithms.compression.pytorch.quantization import QAT_Quantizer\ndummy_input = torch.rand(32, 1, 28, 28).to(device)\nquantizer = QAT_Quantizer(model, config_list, optimizer, dummy_input)\nquantizer.compress()\nfor epoch in range(3):\n    trainer(model, optimizer, criterion)\n    evaluator(model)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "export model and get calibration_config\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "model_path = \"./log/mnist_model.pth\"\ncalibration_path = \"./log/mnist_calibration.pth\"\ncalibration_config = quantizer.export_model(model_path, calibration_path)\n\nprint(\"calibration_config: \", calibration_config)"
+      ]
+    }
+  ],
+  "metadata": {
+    "kernelspec": {
+      "display_name": "Python 3",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "codemirror_mode": {
+        "name": "ipython",
+        "version": 3
+      },
+      "file_extension": ".py",
+      "mimetype": "text/x-python",
+      "name": "python",
+      "nbconvert_exporter": "python",
+      "pygments_lexer": "ipython3",
+      "version": "3.8.8"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
\ No newline at end of file

docs/source/tutorials/quantization_quick_start_mnist.py.md5

+bcaf7880c66acfb20f3e5425730e21de
\ No newline at end of file