pruning_quick_start_mnist.rst 9.27 KB
Newer Older
J-shang's avatar
J-shang committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

.. 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.
25
There are three common practices for pruning a DNN model:
J-shang's avatar
J-shang committed
26

27
28
29
#. Pre-training a model -> Pruning the model -> Fine-tuning the pruned model
#. Pruning a model during training (i.e., pruning aware training) -> Fine-tuning the pruned model
#. Pruning a model -> Training the pruned model from scratch
J-shang's avatar
J-shang committed
30

31
32
NNI supports all of the above pruning practices by working on the key pruning stage.
Following this tutorial for a quick look at how to use NNI to prune a model in a common practice.
J-shang's avatar
J-shang committed
33
34
35
36
37
38

.. GENERATED FROM PYTHON SOURCE LINES 17-22

Preparation
-----------

39
In this tutorial, we use a simple model and pre-trained on MNIST dataset.
J-shang's avatar
J-shang committed
40
41
If you are familiar with defining a model and training in pytorch, you can skip directly to `Pruning Model`_.

42
.. GENERATED FROM PYTHON SOURCE LINES 22-35
J-shang's avatar
J-shang committed
43
44
45
46
47
48
49
50
51
52
53
54
55

.. 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)

56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
    # show the model structure, note that pruner will wrap the model layer.
    print(model)





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

 Out:

 .. code-block:: none

    TorchModel(
      (conv1): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1))
      (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
      (fc1): Linear(in_features=256, out_features=120, bias=True)
      (fc2): Linear(in_features=120, out_features=84, bias=True)
      (fc3): Linear(in_features=84, out_features=10, bias=True)
75
76
77
78
79
80
      (relu1): ReLU()
      (relu2): ReLU()
      (relu3): ReLU()
      (relu4): 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)
81
82
83
84
85
86
87
88
89
90
    )




.. GENERATED FROM PYTHON SOURCE LINES 36-47

.. code-block:: default


J-shang's avatar
J-shang committed
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
    # 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

111
112
113
    Average test loss: 0.5368, Accuracy: 8321/10000 (83%)
    Average test loss: 0.3092, Accuracy: 9104/10000 (91%)
    Average test loss: 0.2070, Accuracy: 9380/10000 (94%)
J-shang's avatar
J-shang committed
114
115
116
117




118
.. GENERATED FROM PYTHON SOURCE LINES 48-58
J-shang's avatar
J-shang committed
119
120
121
122

Pruning Model
-------------

123
124
Using L1NormPruner to prune the model and generate the masks.
Usually, a pruner requires original model and ``config_list`` as its inputs.
125
Detailed about how to write ``config_list`` please refer :doc:`compression config specification <../compression/compression_config_list>`.
J-shang's avatar
J-shang committed
126

127
The following `config_list` means all layers whose type is `Linear` or `Conv2d` will be pruned,
J-shang's avatar
J-shang committed
128
129
130
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.

131
.. GENERATED FROM PYTHON SOURCE LINES 58-67
J-shang's avatar
J-shang committed
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

.. code-block:: default


    config_list = [{
        'sparsity_per_layer': 0.5,
        'op_types': ['Linear', 'Conv2d']
    }, {
        'exclude': True,
        'op_names': ['fc3']
    }]








151
.. GENERATED FROM PYTHON SOURCE LINES 68-69
J-shang's avatar
J-shang committed
152
153
154

Pruners usually require `model` and `config_list` as input arguments.

155
.. GENERATED FROM PYTHON SOURCE LINES 69-76
J-shang's avatar
J-shang committed
156
157
158
159

.. code-block:: default


J-shang's avatar
J-shang committed
160
    from nni.compression.pytorch.pruning import L1NormPruner
J-shang's avatar
J-shang committed
161
    pruner = L1NormPruner(model, config_list)
162
163

    # show the wrapped model structure, `PrunerModuleWrapper` have wrapped the layers that configured in the config_list.
J-shang's avatar
J-shang committed
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
    print(model)





.. 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)
190
191
192
193
194
195
      (relu1): ReLU()
      (relu2): ReLU()
      (relu3): ReLU()
      (relu4): 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)
J-shang's avatar
J-shang committed
196
197
198
199
200
    )




201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
.. GENERATED FROM PYTHON SOURCE LINES 77-84

.. code-block:: default


    # 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

    conv1  sparsity :  0.5
    conv2  sparsity :  0.5
    fc1  sparsity :  0.5
    fc2  sparsity :  0.5




.. GENERATED FROM PYTHON SOURCE LINES 85-88
J-shang's avatar
J-shang committed
231

232
233
Speedup the original model with masks, note that `ModelSpeedup` requires an unwrapped model.
The model becomes smaller after speedup,
J-shang's avatar
J-shang committed
234
235
and reaches a higher sparsity ratio because `ModelSpeedup` will propagate the masks across layers.

236
.. GENERATED FROM PYTHON SOURCE LINES 88-97
J-shang's avatar
J-shang committed
237
238
239
240

.. code-block:: default


241
    # need to unwrap the model, if the model is wrapped before speedup
J-shang's avatar
J-shang committed
242
243
    pruner._unwrap_model()

244
    # speedup the model, for more information about speedup, please refer :doc:`pruning_speedup`.
J-shang's avatar
J-shang committed
245
246
247
248
249
250
251
252
253
254
255
256
257
258
    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

259
260
    aten::log_softmax is not Supported! Please report an issue at https://github.com/microsoft/nni. Thanks~
    Note: .aten::log_softmax.12 does not have corresponding mask inference object
261
    /home/nishang/anaconda3/envs/MCM/lib/python3.9/site-packages/torch/_tensor.py:1013: 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 .grad field to be populated 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. (Triggered internally at  /opt/conda/conda-bld/pytorch_1640811803361/work/build/aten/src/ATen/core/TensorBody.h:417.)
262
      return self._grad
J-shang's avatar
J-shang committed
263
264
265
266




267
.. GENERATED FROM PYTHON SOURCE LINES 98-99
J-shang's avatar
J-shang committed
268

269
the model will become real smaller after speedup
J-shang's avatar
J-shang committed
270

271
.. GENERATED FROM PYTHON SOURCE LINES 99-101
J-shang's avatar
J-shang committed
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292

.. 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)
293
294
295
296
297
298
      (relu1): ReLU()
      (relu2): ReLU()
      (relu3): ReLU()
      (relu4): 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)
J-shang's avatar
J-shang committed
299
300
301
302
303
    )




304
.. GENERATED FROM PYTHON SOURCE LINES 102-106
J-shang's avatar
J-shang committed
305
306
307
308

Fine-tuning Compacted Model
---------------------------
Note that if the model has been sped up, you need to re-initialize a new optimizer for fine-tuning.
309
Because speedup will replace the masked big layers with dense small ones.
J-shang's avatar
J-shang committed
310

311
.. GENERATED FROM PYTHON SOURCE LINES 106-110
J-shang's avatar
J-shang committed
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328

.. code-block:: default


    optimizer = SGD(model.parameters(), 1e-2)
    for epoch in range(3):
        trainer(model, optimizer, criterion)








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

329
   **Total running time of the script:** ( 0 minutes  58.337 seconds)
J-shang's avatar
J-shang committed
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357


.. _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>`_