Posts on Decsnow Media

MIT 6.5940 EfficientML.ai Lab 3: Neural Architecture Search

Tue, 02 Jun 2026 00:00:00 +0000

MIT 6.5940 EfficientML.ai Lab 3: Neural Architecture Search

by MIT HAN Lab

Introduction

This colab notebook provides code and a framework for Lab 3: neural architecture search. In this lab, you will learn how to search for a tiny neural network that can run efficiently on a microcontroller. You can work out your solutions here.

Over a long period of time, researchers manually design neural network architectures. The design space of NN architectures is very large: it includes #layers, #channel width, #branches, kernel sizes and input resolutions. As a result, tuning these design knobs manually is notoriously hard. Neural architecture search, or NAS, on the other hand, can help researchers automatically tune these design knobs under various type of efficiency and accuracy constraints. As a result, it greatly saves engineering cost in NN design and helps democratize AI. In this lab, we are going to walk you through neural architecture search from scratch.

Early NAS methods train candidate network in the design space exhaustively and use RNN-based controllers with reinforcement learning to optimize the sampling strategy. The representative methods include Neural Architecture Search with Reinforcement Learning, NASNet and MNASNet. These methods are usually very computationally expensive because each candidate network has to be trained from scratch such that the RNN-based controller can get the reward signal (which is the accuracy of the candidate network).

Later, researchers develop differentiable NAS methods such as DARTS, ProxylessNAS and FBNet that greatly reduces the total cost of training candidate networks. DARTS model the output of each layer as a weighted average of outputs from different candidate operations, and ProxylessNAS further reduces the memory cost of DARTS by keeping only two paths instead of all paths in the memory. Later one-shot methods such as Single Path One Shot further notice that it is possible to keep only one path at a time during the training process.

Albeit being much more efficient than controller-based methods, differentiable NAS and one-shot NAS still requires running the entire training, search and finetune pipeline everytime we design a new neural network. This brings about large cost (typically 200-300 GPU hours for the ImageNet dataset) for model specialization considering the large amount of edge devices (e.g. there are > 20 billion IoT devices till 2018).

In this lab, we therefore refer to Once for All (OFA), a method that can greatly reduce the cost of specialize NN architectures for different devices. OFA trains a large super network that contains all sub-networks within the design space. If we directly extract the sub-networks from the super network, they can achieve similar-level of accuracy compared with training from scratch. As such, OFA supports direct deployment with no retrain.

Furthermore, OFA introduces accuracy and efficiency predictors to further reduce the evaluation cost during architecture search. Intuitively the accuracy of a sub-network requires running inference on the entire holdout validation set, which can take around 1 minute on ImageNet. OFA, instead, collects a large amount of (architecture, accuracy) pairs beforehand and trains a regression model to predict the accuracy during search. This greatly reduces the cost to get the accuracy feedback from 1 minute to less than 1 second for each sub-network. Similar idea can also be applied to efficiency predictors, where the evaluation of latency are usually very slow since we have to run the forward pass of the candidate network for many times.

In this lab, you will study how to search for efficient networks that can run on extremely resource-constrained microcontrollers with OFA and predictors. Microcontrollers are low-cost, low-power hardware. They are widely deployed and have wide applications.

But the tight memory budget (50,000x smaller than GPUs) makes deep learning deployment difficult.

There are 2 main sections: accuracy & efficiency predictors and architecture search.

For predictors, there are 4 questions in total. There is one question (5 pts) in the Getting Started section and the other three questions (30 pts) are in the Predictors section.
For architecture search, there are 6 questions in total.

First, install the required packages and download the Visual Wake Words dataset that will be used in this lab.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15


# print("Cleanning up workspace ...")
# !rm -rf *
print("Installing graphviz ...")
!sudo apt-get install graphviz 1>/dev/null
print("Downloading MCUNet codebase ...")
!wget https://www.dropbox.com/s/3y2n2u3mfxczwcb/mcunetv2-dev-main.zip?dl=0 >/dev/null
!unzip mcunetv2-dev-main.zip* 1>/dev/null
!mv mcunetv2-dev-main/* . 1>/dev/null
print("Downloading VWW dataset ...")
!wget https://www.dropbox.com/s/169okcuuv64d4nn/data.zip?dl=0 >/dev/null
print("Unzipping VWW dataset ...")
!unzip data.zip* 1>/dev/null
print("Installing thop and onnx ...")
!pip install thop 1>/dev/null
!pip install onnx 1>/dev/null

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53


Installing graphviz ...
Downloading MCUNet codebase ...
--2026-05-29 07:13:12-- https://www.dropbox.com/s/3y2n2u3mfxczwcb/mcunetv2-dev-main.zip?dl=0
Resolving www.dropbox.com (www.dropbox.com)... 162.125.81.18, 2620:100:6030:18::a27d:5012
Connecting to www.dropbox.com (www.dropbox.com)|162.125.81.18|:443... connected.
HTTP request sent, awaiting response... 302 Found
Location: https://www.dropbox.com/scl/fi/4fo3qexvbdg47cb7e1mkg/mcunetv2-dev-main.zip?rlkey=241ldxul5ymto3ins60bs3x2m&dl=0 [following]
--2026-05-29 07:13:13-- https://www.dropbox.com/scl/fi/4fo3qexvbdg47cb7e1mkg/mcunetv2-dev-main.zip?rlkey=241ldxul5ymto3ins60bs3x2m&dl=0
Reusing existing connection to www.dropbox.com:443.
HTTP request sent, awaiting response... 302 Found
Location: https://uce2ffe3a4112d0e3bee439b15e3.dl.dropboxusercontent.com/cd/0/inline/DBUpMt5LAfioMYnE8NXg4AXxaOAgP3IRSaLfgQJlKGYLvBZBML4ymKbVGl9g6trf4-p2104PVFNCn2bRiznZvhwMG_l_Ypg8NYEqRA0Vd4QhEePYLD0bCSWWhfWwnWUHSrOrOt9I4Xj7N4mJupaEi_b8/file# [following]
--2026-05-29 07:13:14-- https://uce2ffe3a4112d0e3bee439b15e3.dl.dropboxusercontent.com/cd/0/inline/DBUpMt5LAfioMYnE8NXg4AXxaOAgP3IRSaLfgQJlKGYLvBZBML4ymKbVGl9g6trf4-p2104PVFNCn2bRiznZvhwMG_l_Ypg8NYEqRA0Vd4QhEePYLD0bCSWWhfWwnWUHSrOrOt9I4Xj7N4mJupaEi_b8/file
Resolving uce2ffe3a4112d0e3bee439b15e3.dl.dropboxusercontent.com (uce2ffe3a4112d0e3bee439b15e3.dl.dropboxusercontent.com)... 162.125.81.15, 2620:100:6030:15::a27d:500f
Connecting to uce2ffe3a4112d0e3bee439b15e3.dl.dropboxusercontent.com (uce2ffe3a4112d0e3bee439b15e3.dl.dropboxusercontent.com)|162.125.81.15|:443... connected.
HTTP request sent, awaiting response... 302 Found
Location: /cd/0/inline2/DBWl89d8NUNgIONfMj-zs_ivswFzPO9P6kl_UAT5BZ1s3CP9-I0NW-lmMX68toRfrgt6n6WY8DbROlodXDXvCaWr97zrvD8jZ9MsAktLLGagpqxIspS5KcaR2RKkITVJqf_WNhasuwSYrX5K1dwrlU9EkxlSBSp2JlqdoB85RqdOYTdnBtFRsDlzsyMPOz96B75Jlk5MiqcxYPRDYb24dpQbaOshcpI_xZW2hlNEM1gHklA8-17QxsCfKkjJCMctkp_Wv72BG_R2Lo5NRuu3ZVE3wMq1uGpdBAQhOBx8KCwrxrXvwgXfCcOiyDDqtBuflqUoG0mveeMJe1gxgTaZQAxR9j0bzItDUMkm87WQL0I2ls5bZ014NPWD9I69muPxy4w/file [following]
--2026-05-29 07:13:15-- https://uce2ffe3a4112d0e3bee439b15e3.dl.dropboxusercontent.com/cd/0/inline2/DBWl89d8NUNgIONfMj-zs_ivswFzPO9P6kl_UAT5BZ1s3CP9-I0NW-lmMX68toRfrgt6n6WY8DbROlodXDXvCaWr97zrvD8jZ9MsAktLLGagpqxIspS5KcaR2RKkITVJqf_WNhasuwSYrX5K1dwrlU9EkxlSBSp2JlqdoB85RqdOYTdnBtFRsDlzsyMPOz96B75Jlk5MiqcxYPRDYb24dpQbaOshcpI_xZW2hlNEM1gHklA8-17QxsCfKkjJCMctkp_Wv72BG_R2Lo5NRuu3ZVE3wMq1uGpdBAQhOBx8KCwrxrXvwgXfCcOiyDDqtBuflqUoG0mveeMJe1gxgTaZQAxR9j0bzItDUMkm87WQL0I2ls5bZ014NPWD9I69muPxy4w/file
Reusing existing connection to uce2ffe3a4112d0e3bee439b15e3.dl.dropboxusercontent.com:443.
HTTP request sent, awaiting response... 200 OK
Length: 31675052 (30M) [application/zip]
Saving to: ‘mcunetv2-dev-main.zip?dl=0’

mcunetv2-dev-main.z 100%[===================>] 30.21M 13.6MB/s in 2.2s 

2026-05-29 07:13:17 (13.6 MB/s) - ‘mcunetv2-dev-main.zip?dl=0’ saved [31675052/31675052]

Downloading VWW dataset ...
--2026-05-29 07:13:18-- https://www.dropbox.com/s/169okcuuv64d4nn/data.zip?dl=0
Resolving www.dropbox.com (www.dropbox.com)... 162.125.81.18, 2620:100:6030:18::a27d:5012
Connecting to www.dropbox.com (www.dropbox.com)|162.125.81.18|:443... connected.
HTTP request sent, awaiting response... 302 Found
Location: https://www.dropbox.com/scl/fi/a2d3mbhj0tn0dfhnojkr6/data.zip?rlkey=oo4sfh74waerp9ub0lecy8u7d&dl=0 [following]
--2026-05-29 07:13:19-- https://www.dropbox.com/scl/fi/a2d3mbhj0tn0dfhnojkr6/data.zip?rlkey=oo4sfh74waerp9ub0lecy8u7d&dl=0
Reusing existing connection to www.dropbox.com:443.
HTTP request sent, awaiting response... 302 Found
Location: https://uc43f5aa2e40d8baeff72b586e14.dl.dropboxusercontent.com/cd/0/inline/DBVzoQ0EEZw3oWOjbMaxjvmfMlusYlUIafr2zCkAaSLgzOL8f6lJcRq1iik3hJdj8QFFIRas8xkm1yaOOOHnXPAXTRqElbGJXEgpNDpMuhK7W67d6XS2J7r14wHjhR3qCTNsn8BKNqaBrA8PDk7-xF2T/file# [following]
--2026-05-29 07:13:20-- https://uc43f5aa2e40d8baeff72b586e14.dl.dropboxusercontent.com/cd/0/inline/DBVzoQ0EEZw3oWOjbMaxjvmfMlusYlUIafr2zCkAaSLgzOL8f6lJcRq1iik3hJdj8QFFIRas8xkm1yaOOOHnXPAXTRqElbGJXEgpNDpMuhK7W67d6XS2J7r14wHjhR3qCTNsn8BKNqaBrA8PDk7-xF2T/file
Resolving uc43f5aa2e40d8baeff72b586e14.dl.dropboxusercontent.com (uc43f5aa2e40d8baeff72b586e14.dl.dropboxusercontent.com)... 162.125.69.15, 2620:100:6025:15::a27d:450f
Connecting to uc43f5aa2e40d8baeff72b586e14.dl.dropboxusercontent.com (uc43f5aa2e40d8baeff72b586e14.dl.dropboxusercontent.com)|162.125.69.15|:443... connected.
HTTP request sent, awaiting response... 302 Found
Location: /cd/0/inline2/DBUjuskqW4eHT9cQWv6pT8QMYekVsvYETcRxWO87cyvlt9azKUNt1ddj6WYhlGjvatnejN35om-nmMArddZeDEo2-UF-lqXN2oeBgXV2OlXC-HyXi7IxhLri7VtG_xfRIgmzcVr8UU1nRqssBmwsDTjgwZG-xVt6PPQlItoZ4OKVIirEsncNP_2yJ9LtNXQjr4VBP5ylt5rEpvymX3bxkrebky3nxo8oJtR92WRc9Kzusz826DsoP6jwR8MoX__ibfMnmJU80RLuovnBfNlbfWpSlhJ5NrVxsrx-heecrSfYNh6PE1r-HsvYfXf6id2PCkTnxDLOqQt7KxrXAG3reLr4VdJi5fdY7p6gJQgH1LRLlBzbnbPBhgV4nKCvoe227Ww/file [following]
--2026-05-29 07:13:21-- https://uc43f5aa2e40d8baeff72b586e14.dl.dropboxusercontent.com/cd/0/inline2/DBUjuskqW4eHT9cQWv6pT8QMYekVsvYETcRxWO87cyvlt9azKUNt1ddj6WYhlGjvatnejN35om-nmMArddZeDEo2-UF-lqXN2oeBgXV2OlXC-HyXi7IxhLri7VtG_xfRIgmzcVr8UU1nRqssBmwsDTjgwZG-xVt6PPQlItoZ4OKVIirEsncNP_2yJ9LtNXQjr4VBP5ylt5rEpvymX3bxkrebky3nxo8oJtR92WRc9Kzusz826DsoP6jwR8MoX__ibfMnmJU80RLuovnBfNlbfWpSlhJ5NrVxsrx-heecrSfYNh6PE1r-HsvYfXf6id2PCkTnxDLOqQt7KxrXAG3reLr4VdJi5fdY7p6gJQgH1LRLlBzbnbPBhgV4nKCvoe227Ww/file
Reusing existing connection to uc43f5aa2e40d8baeff72b586e14.dl.dropboxusercontent.com:443.
HTTP request sent, awaiting response... 200 OK
Length: 375960448 (359M) [application/zip]
Saving to: ‘data.zip?dl=0’

data.zip?dl=0 100%[===================>] 358.54M 15.7MB/s in 25s 

2026-05-29 07:13:46 (14.5 MB/s) - ‘data.zip?dl=0’ saved [375960448/375960448]

Unzipping VWW dataset ...
Installing thop and onnx ...

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26


import argparse
import json
from PIL import Image
from tqdm import tqdm
import copy
import math
import numpy as np
import os
import random
import torch
from torch import nn
from torchvision import datasets, transforms
from mcunet.tinynas.search.accuracy_predictor import (
 AccuracyDataset,
 MCUNetArchEncoder,
)

from mcunet.tinynas.elastic_nn.networks.ofa_mcunets import OFAMCUNets
from mcunet.utils.mcunet_eval_helper import calib_bn, validate
from mcunet.utils.arch_visualization_helper import draw_arch


%matplotlib inline
from matplotlib import pyplot as plt
import warnings
warnings.filterwarnings('ignore')

Getting Started: Super Network and the VWW dataset (1 Question, 5 pts)

In this lab, we will be using the MCUNetV2 super network trained in an once-for-all (OFA) manner. Recall that super network is a randomized large neural network that contains all candidate subnets within the design space. We can directly extract the subnets from the super network and evaluate their accuracy. The accuracy can be further used as a feedback signal to guide neural network design. The advantage of OFA super network is that the directly extracted subnets can achieve similar (or even better) performance compared with training from scratch.

MCUNetV2 is a family of efficiency neural networks tailored for resource-constrained microntrollers. It utilizes patch-based inference, receptive field redistribution and system-NN co-design and greatly improves the accuracy-efficiency tradeoff of MCUNet.

Let’s first visualize some samples in the VWW dataset. This is a binary image classficiation (whether people is present in the image) dataset subsampled from Microsoft COCO. We first define a function to set up a dataloader over the validation set.

Note: The function build_val_data_loader has an argument split. We use split = 0 (default value) to represent the validation set (cannot be directly used for architecture search), and split = 1 will be used as a holdout minival set (used to generate the accuracy dataset and calibrate BN parameters).

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25


def build_val_data_loader(data_dir, resolution, batch_size=128, split=0):
 # split = 0: real val set, split = 1: holdout validation set
 assert split in [0, 1]
 normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
 kwargs = {"num_workers": min(8, os.cpu_count()), "pin_memory": False}

 val_transform = transforms.Compose(
 [
 transforms.Resize(
 (resolution, resolution)
 ), # if center crop, the person might be excluded
 transforms.ToTensor(),
 normalize,
 ]
 )
 val_dataset = datasets.ImageFolder(data_dir, transform=val_transform)

 val_dataset = torch.utils.data.Subset(
 val_dataset, list(range(len(val_dataset)))[split::2]
 )

 val_loader = torch.utils.data.DataLoader(
 val_dataset, batch_size=batch_size, shuffle=False, **kwargs
 )
 return val_loader

Using that dataloader builder, we are able to navigate through the VWW validation set. You can run the following cell for several times to see different images in the dataset.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24


data_dir = "data/vww-s256/val"

val_data_loader = build_val_data_loader(data_dir, resolution=128, batch_size=1)

vis_x, vis_y = 2, 3
fig, axs = plt.subplots(vis_x, vis_y)

num_images = 0
for data, label in val_data_loader:
 img = np.array((((data + 1) / 2) * 255).numpy(), dtype=np.uint8)
 img = img[0].transpose(1, 2, 0)
 if label.item() == 0:
 label_text = "No person"
 else:
 label_text = "Person"
 axs[num_images // vis_y][num_images % vis_y].imshow(img)
 axs[num_images // vis_y][num_images % vis_y].set_title(f"Label: {label_text}")
 axs[num_images // vis_y][num_images % vis_y].set_xticks([])
 axs[num_images // vis_y][num_images % vis_y].set_yticks([])
 num_images += 1
 if num_images > vis_x * vis_y - 1:
 break

plt.show()

1

<Figure size 640x480 with 6 Axes>

Cool, now you have a basic idea about the dataset. Let’s then construct the OFA super network! The OFAMCUNets super network is composed of $>10^{19}$ subnets in the MCUNetV2 design space. The subnets are composed of inverted MobileNet blocks with different kernel sizes (3, 5, 7) and expand ratios (3, 4, 6). The OFA super network also allows elastic depths (base depth to base_depth + 2) for all network stages. Finally, the super network supports global channel scaling (specified by width_mult_list) by 0.5$\times$, 0.75$\times$ or 1.0$\times$.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20


device = "cuda:0"
ofa_network = OFAMCUNets(
 n_classes=2,
 bn_param=(0.1, 1e-3),
 dropout_rate=0.0,
 base_stage_width="mcunet384",
 width_mult_list=[0.5, 0.75, 1.0],
 ks_list=[3, 5, 7],
 expand_ratio_list=[3, 4, 6],
 depth_list=[0, 1, 2],
 base_depth=[1, 2, 2, 2, 2],
 fuse_blk1=True,
 se_stages=[False, [False, True, True, True], True, True, True, False],
)

ofa_network.load_state_dict(
 torch.load("vww_supernet.pth", map_location="cpu")["state_dict"], strict=True
)

ofa_network = ofa_network.to(device)

We then verify that the checkpoint is correctly loaded. We will sample some networks in the MCUNetV2 design space and evaluate its accuracy on the VWW dataset. The evaluation will take less than one minutes, and you are expected to see an accuracy around 83.6-88.7%. As you can see, we can directly extract these subnets from the design space and get their accuracy very quickly without training. This is a unique advantage brought by once-for-all (OFA) super networks.

Let’s first define a helper function evaluate_sub_network that testes the accuracy of a sub network directly extracted from the super network.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21


from mcunet.utils.pytorch_utils import count_peak_activation_size, count_net_flops, count_parameters

def evaluate_sub_network(ofa_network, cfg, image_size=None):
 if "image_size" in cfg:
 image_size = cfg["image_size"]
 batch_size = 128
 # step 1. sample the active subnet with the given config.
 ofa_network.set_active_subnet(**cfg)
 # step 2. extract the subnet with corresponding weights.
 subnet = ofa_network.get_active_subnet().to(device)
 # step 3. calculate the efficiency stats of the subnet.
 peak_memory = count_peak_activation_size(subnet, (1, 3, image_size, image_size))
 macs = count_net_flops(subnet, (1, 3, image_size, image_size))
 params = count_parameters(subnet)
 # step 4. perform BN parameter re-calibration.
 calib_bn(subnet, data_dir, batch_size, image_size)
 # step 5. define the validation dataloader.
 val_loader = build_val_data_loader(data_dir, image_size, batch_size)
 # step 6. validate the accuracy.
 acc = validate(subnet, val_loader)
 return acc, peak_memory, macs, params

We also provide a handly helper function to visualize the architecture of the subnets. The function takes in the configuration of the subnet and returns an image representing the architecture.

1
2
3
4
5
6
7
8


def visualize_subnet(cfg):
 draw_arch(cfg["ks"], cfg["e"], cfg["d"], cfg["image_size"], out_name="viz/subnet")
 im = Image.open("viz/subnet.png")
 im = im.rotate(90, expand=1)
 fig = plt.figure(figsize=(im.size[0] / 250, im.size[1] / 250))
 plt.axis("off")
 plt.imshow(im)
 plt.show()

Now, let’s visualize some subnets and evaluate them on the VWW dataset! We provide an example to randomly sample a subnet from the design space, and get its accuracy, macs, parameters on the VWW dataset. We also visualize the architecture using visualize_subnet.

In the architecture visualization, the legend of each block MBConv{e}-{k}x{k} means that the current block is a mobile inverted block with expand ratio e and the kernel size of the depthwise convolution layer is k. Different colors of the blocks indicate different kernel sizes, and gray blocks are network stage dividers. Different widths for the blocks indicate different expand ratios. We also annotate the output resolution close to each block.

Note that we assume that the image resolution is fixed to be 96. Feel free to add another cell below and play with the input resolution.

Hint: you can change the sample_function argument of the sample_active_subnet method to control the sampling process.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16


image_size = 96

cfg = ofa_network.sample_active_subnet(sample_function=random.choice, image_size=image_size)
acc, _, _, params = evaluate_sub_network(ofa_network, cfg)
visualize_subnet(cfg)
print(f"The accuracy of the sampled subnet: #params={params/1e6: .1f}M, accuracy={acc: .1f}%.")

largest_cfg = ofa_network.sample_active_subnet(sample_function=max, image_size=image_size)
acc, _, _, params = evaluate_sub_network(ofa_network, largest_cfg)
visualize_subnet(largest_cfg)
print(f"The largest subnet: #params={params/1e6: .1f}M, accuracy={acc: .1f}%.")

smallest_cfg = ofa_network.sample_active_subnet(sample_function=min, image_size=image_size)
acc, peak_memory, macs, params = evaluate_sub_network(ofa_network, smallest_cfg)
visualize_subnet(smallest_cfg)
print(f"The smallest subnet: #params={params/1e6: .1f}M, accuracy={acc: .1f}%.")

1

Validate: 100%|██████████| 32/32 [00:08<00:00, 3.75it/s, loss=0.325, top1=86.6]

1

<Figure size 1491.6x177.2 with 1 Axes>

1

The accuracy of the sampled subnet: #params= 0.6M, accuracy= 86.6%.

1

Validate: 100%|██████████| 32/32 [00:07<00:00, 4.52it/s, loss=0.29, top1=88.6]

1

<Figure size 1769.2x177.2 with 1 Axes>

1

The largest subnet: #params= 2.5M, accuracy= 88.6%.

1

Validate: 100%|██████████| 32/32 [00:08<00:00, 3.96it/s, loss=0.379, top1=83.4]

1

<Figure size 1075.6x177.2 with 1 Axes>

1

The smallest subnet: #params= 0.3M, accuracy= 83.4%.

Question 1 (5 pts): Design space exploration.

Try manually sample different subnets by running the cell above multiple times. You can also vary the input resolution. Talk about your findings.

Hint: which dimension plays the most important role for the accuracy?

Answer: At 96×96 resolution, a random subnet achieved ~86.6% accuracy, the largest subnet ~88.6%, and the smallest ~83.4%. Architecture choices (width, depth, kernel size, expand ratio) still cause a noticeable spread (~5 points), but varying image_size has the largest effect: higher resolution generally improves accuracy at the cost of more MACs and peak activation memory. Among architecture knobs, input image resolution plays the most important role for accuracy.

Part 1. Predictors (3 Questions, 30 pts)

Neural architecture search requires sampling a large amount of sub-networks from the OFA supernet and evaluate the performance of these sub-networks. Such performance evaluation is time-consuming.

In this lab, we explore very fast neural network search with efficiency predictors and accuracy predictors.

Question 2 (10 pts): Implement the efficiency predictor.

For the efficiency predictor, we use an hook-based analytical model to count the #MACs and peak memory consumption of a given network. Let’s build it from scratch using our provided APIs.

Specifically, we define a class called AnalyticalEfficiencyPredictor. There are two major functions in this class, get_efficiency and satisfy_constraint.

The function get_efficiency takes in the subnet configuration and returns the #MACs and peak memory of the given subnet. Here, we assume the unit for the #MACs is million and the unit of the peak memory consumption is KB.

Hint: take a look at the evaluate_sub_network function above. Let’s use count_net_flops to get the MACs of the network and count_peak_activation_size to get the activation size of the network.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29


class AnalyticalEfficiencyPredictor:
 def __init__(self, net):
 self.net = net

 def get_efficiency(self, spec: dict):
 self.net.set_active_subnet(**spec)
 subnet = self.net.get_active_subnet()
 if torch.cuda.is_available():
 subnet = subnet.cuda()
 ############### YOUR CODE STARTS HERE ###############
 # Hint: take a look at the `evaluate_sub_network` function above.
 # Hint: the data shape is (batch_size, input_channel, image_size, image_size)
 data_shape = (1, 3, spec["image_size"], spec["image_size"])
 macs = count_net_flops(subnet, data_shape)
 peak_memory = count_peak_activation_size(subnet, data_shape)
 ################ YOUR CODE ENDS HERE ################

 return dict(millionMACs=macs / 1e6, KBPeakMemory=peak_memory / 1024)

 def satisfy_constraint(self, measured: dict, target: dict):
 for key in measured:
 # if the constraint is not specified, we just continue
 if key not in target:
 continue
 # if we exceed the constraint, just return false.
 if measured[key] > target[key]:
 return False
 # no constraint violated, return true.
 return True

Let’s test your implementation for the analytical efficiency predictor by examining the returned values for the smallest and largest subnets we just evaluated a while ago. The results from the efficiency predictor should match with the previous results.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


efficiency_predictor = AnalyticalEfficiencyPredictor(ofa_network)

image_size = 96
# Print out the efficiency of the smallest subnet.
smallest_cfg = ofa_network.sample_active_subnet(sample_function=min, image_size=image_size)
eff_smallest = efficiency_predictor.get_efficiency(smallest_cfg)

# Print out the efficiency of the largest subnet.
largest_cfg = ofa_network.sample_active_subnet(sample_function=max, image_size=image_size)
eff_largest = efficiency_predictor.get_efficiency(largest_cfg)

print("Efficiency stats of the smallest subnet:", eff_smallest)
print("Efficiency stats of the largest subnet:", eff_largest)

1
2


Efficiency stats of the smallest subnet: {'millionMACs': 8.302128, 'KBPeakMemory': 72.0}
Efficiency stats of the largest subnet: {'millionMACs': 79.416432, 'KBPeakMemory': 270.0}

Question 3 (10 pts): Implement the accuracy predictor.

For the accuracy predictor, it predicts the classification accuracy of a given sub-network on the VWW dataset so that we do NOT need to run costly inference every time when we encounter a new subnet during architecture search. Such an accuracy predictor is an MLP (multi-layer perception) model trained on an accuracy dataset built with the OFA network. The inference of an MLP network takes only a few milliseconds, thus the accuracy predictor can speedup the search process by orders of magnitude.

The accuracy predictor takes in the architecture of a sub-network and predicts its accuracy on the VWW dataset. Since it is an MLP network, the sub-network must be encoded into a vector. In this lab, we provide a class MCUNetArchEncoder to perform such conversion from sub-network architecture to a binary vector.

1
2
3
4
5
6
7
8


image_size_list = [96, 112, 128, 144, 160]
arch_encoder = MCUNetArchEncoder(
 image_size_list=image_size_list,
 base_depth=ofa_network.base_depth,
 depth_list=ofa_network.depth_list,
 expand_list=ofa_network.expand_ratio_list,
 width_mult_list=ofa_network.width_mult_list,
)

We generated an accuracy dataset beforehand, which is a collection of [architecture, accuracy] pairs stored under the acc_datasets folder.

With the architecture encoder, you are now required define the accuracy predictor, which is a multi-layer perception (MLP) network with 400 channels per intermediate layer. For simplicity, we fix the number of layers to be 3. Please implement this MLP network in the following cell.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58


class AccuracyPredictor(nn.Module):
 def __init__(
 self,
 arch_encoder,
 hidden_size=400,
 n_layers=3,
 checkpoint_path=None,
 device="cuda:0",
 ):
 super(AccuracyPredictor, self).__init__()
 self.arch_encoder = arch_encoder
 self.hidden_size = hidden_size
 self.n_layers = n_layers
 self.device = device

 layers = []

 ############### YOUR CODE STARTS HERE ###############
 # Let's build an MLP with n_layers layers.
 # Each layer (nn.Linear) has hidden_size channels and
 # uses nn.ReLU as the activation function.
 # Hint: You can assume that n_layers is fixed to be 3, for simplicity.
 # Hint: the input dimension of the first layer is not hidden_size.
 # use self.arch_encoder.n_dim to get the input dimension
 for i in range(self.n_layers):
 layers.append(
 nn.Sequential(
 nn.Linear(
 self.arch_encoder.n_dim if i == 0 else self.hidden_size,
 self.hidden_size,
 ),
 nn.ReLU(inplace=True),
 )
 )
 ################ YOUR CODE ENDS HERE ################
 layers.append(nn.Linear(self.hidden_size, 1, bias=False))
 self.layers = nn.Sequential(*layers)
 self.base_acc = nn.Parameter(
 torch.zeros(1, device=self.device), requires_grad=False
 )

 if checkpoint_path is not None and os.path.exists(checkpoint_path):
 checkpoint = torch.load(checkpoint_path, map_location="cpu")
 if "state_dict" in checkpoint:
 checkpoint = checkpoint["state_dict"]
 self.load_state_dict(checkpoint)
 print("Loaded checkpoint from %s" % checkpoint_path)

 self.layers = self.layers.to(self.device)

 def forward(self, x):
 y = self.layers(x).squeeze()
 return y + self.base_acc

 def predict_acc(self, arch_dict_list):
 X = [self.arch_encoder.arch2feature(arch_dict) for arch_dict in arch_dict_list]
 X = torch.tensor(np.array(X)).float().to(self.device)
 return self.forward(X)

Let’s print out the architecture of the AccuracyPredictor you just defined.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12


os.makedirs("pretrained", exist_ok=True)
acc_pred_checkpoint_path = (
 f"pretrained/{ofa_network.__class__.__name__}_acc_predictor.pth"
)
acc_predictor = AccuracyPredictor(
 arch_encoder,
 hidden_size=400,
 n_layers=3,
 checkpoint_path=None,
 device=device,
)
print(acc_predictor)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17


AccuracyPredictor(
 (layers): Sequential(
 (0): Sequential(
 (0): Linear(in_features=128, out_features=400, bias=True)
 (1): ReLU(inplace=True)
 )
 (1): Sequential(
 (0): Linear(in_features=400, out_features=400, bias=True)
 (1): ReLU(inplace=True)
 )
 (2): Sequential(
 (0): Linear(in_features=400, out_features=400, bias=True)
 (1): ReLU(inplace=True)
 )
 (3): Linear(in_features=400, out_features=1, bias=False)
 )
)

Let’s first visualize some samples in the accuracy dataset in the following cell.

The accuracy dataset is composed of 50,000 [architecture, accuracy] pairs, where 40,000 of them are used as the training set and the rest 10,000 are used as validation set.

For accuracy, We calculate the average accuracy of all [architecture, accuracy] pairs on the accuracy dataset and define it as base_acc. For the accuracy predictor, instead of directly regressing the accuracy of each architecture, its training target is accuracy - base_acc. Since accuracy - base_acc is usually much smaller than accuracy itself, this can make training easier.

For architecture, each subnet within the design space is uniquely represented by a binary vector. The binary vector is a concatenation of the one-hot representation for both global parameters (e.g. input resolution, width multiplier) and parameters of each inverted MobileNet block (e.g. kernel sizes and expand ratios). Note that we prefer one-hot representations over numerical representations because all design hyperparameters are discrete values.

For example, our design space supports

1
2


kernel_size = [3, 5, 7]
expand_ratio = [3, 4, 6]

Then, we represent kernel_size=3 as [1, 0, 0], kernel_size=5 as [0, 1, 0], and kernel_size=7 as [0, 0, 1]. Similarly, for expand_ratio=3, it is written as [1, 0, 0]; expand_ratio=4 is written as [0, 1, 0] and expand_ratio=6 is written as [0, 0, 1]. The representation for each inverted MobileNet block is obtained by concatenating the kernel size embedding with the expand ratio embedding. Note that for skipped blocks, we use [0, 0, 0] to represent their kernel sizes and expand ratios. You will see a detailed explanation of the architecture-embedding correspondence after running the following cell.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22


acc_dataset = AccuracyDataset("acc_datasets")
train_loader, valid_loader, base_acc = acc_dataset.build_acc_data_loader(
 arch_encoder=arch_encoder
)

print(f"The basic accuracy (mean accuracy of all subnets within the dataset is: {(base_acc * 100): .1f}%.")

# Let's print one sample in the training set
sampled = 0
for (data, label) in train_loader:
 data = data.to(device)
 label = label.to(device)
 print("=" * 100)
 # dummy pass to print the divided encoding
 arch_encoding = arch_encoder.feature2arch(data[0].int().cpu().numpy(), verbose=False)
 # print out the architecture encoding process in detail
 arch_encoding = arch_encoder.feature2arch(data[0].int().cpu().numpy(), verbose=True)
 visualize_subnet(arch_encoding)
 print(f"The accuracy of this subnet on the holdout validation set is: {(label[0] * 100): .1f}%.")
 sampled += 1
 if sampled == 1:
 break

1

Loading data: 100%|██████████| 50000/50000 [00:00<00:00, 84005.60it/s]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32


Train Size: 40000, Valid Size: 10000
The basic accuracy (mean accuracy of all subnets within the dataset is: 90.3%.
====================================================================================================
network embedding: [0 1 0 0 0 | 1 0 0 | 1 0 0 | 0 0 1 | 1 0 0 | 0 1 0 | 0 1 0 | 1 0 0 | 0 0 1 | 0 0 1 | 0 1 0 | 0 0 1 | 0 0 1 | 0 1 0 | 1 0 0 | 0 1 0 | 1 0 0 | 1 0 0 | 1 0 0 | 1 0 0 | 0 0 0 | 0 0 0 | 0 0 0 | 0 0 0 | 0 0 1 | 0 1 0 | 1 0 0 | 0 0 1 | 0 0 0 | 0 0 0 | 0 0 0 | 0 0 0 | 0 1 0 | 0 1 0 | 0 1 0 | 0 1 0 | 0 0 1 | 0 1 0 | 0 0 0 | 0 0 0 | 0 0 1 | 1 0 0]
image resolution embedding: [0 1 0 0 0] => image resolution: 112
width multiplier embedding: [1 0 0] => width multiplier: 0.5
**************************************************Stage1**************************************************
kernel size embedding: [1 0 0] => kernel size: 3; expand ratio embedding: [0 0 1] => expand ratio: 6
kernel size embedding: [1 0 0] => kernel size: 3; expand ratio embedding: [0 1 0] => expand ratio: 4
kernel size embedding: [0 1 0] => kernel size: 5; expand ratio embedding: [1 0 0] => expand ratio: 3
**************************************************Stage2**************************************************
kernel size embedding: [0 0 1] => kernel size: 7; expand ratio embedding: [0 0 1] => expand ratio: 6
kernel size embedding: [0 1 0] => kernel size: 5; expand ratio embedding: [0 0 1] => expand ratio: 6
kernel size embedding: [0 0 1] => kernel size: 7; expand ratio embedding: [0 1 0] => expand ratio: 4
kernel size embedding: [1 0 0] => kernel size: 3; expand ratio embedding: [0 1 0] => expand ratio: 4
**************************************************Stage3**************************************************
kernel size embedding: [1 0 0] => kernel size: 3; expand ratio embedding: [1 0 0] => expand ratio: 3
kernel size embedding: [1 0 0] => kernel size: 3; expand ratio embedding: [1 0 0] => expand ratio: 3
kernel size embedding: [0 0 0] expand ratio embedding: [0 0 0] => layer skipped.
kernel size embedding: [0 0 0] expand ratio embedding: [0 0 0] => layer skipped.
**************************************************Stage4**************************************************
kernel size embedding: [0 0 1] => kernel size: 7; expand ratio embedding: [0 1 0] => expand ratio: 4
kernel size embedding: [1 0 0] => kernel size: 3; expand ratio embedding: [0 0 1] => expand ratio: 6
kernel size embedding: [0 0 0] expand ratio embedding: [0 0 0] => layer skipped.
kernel size embedding: [0 0 0] expand ratio embedding: [0 0 0] => layer skipped.
**************************************************Stage5**************************************************
kernel size embedding: [0 1 0] => kernel size: 5; expand ratio embedding: [0 1 0] => expand ratio: 4
kernel size embedding: [0 1 0] => kernel size: 5; expand ratio embedding: [0 1 0] => expand ratio: 4
kernel size embedding: [0 0 1] => kernel size: 7; expand ratio embedding: [0 1 0] => expand ratio: 4
kernel size embedding: [0 0 0] expand ratio embedding: [0 0 0] => layer skipped.
**************************************************Stage6**************************************************
kernel size embedding: [0 0 1] => kernel size: 7; expand ratio embedding: [1 0 0] => expand ratio: 3

1

<Figure size 1422.4x177.2 with 1 Axes>

1

The accuracy of this subnet on the holdout validation set is: 88.7%.

Question 4 (10 pts): Complete the code for accuracy predictor training.

Now let’s train the accuracy predictor using the dataset we provided! In this part, you are responsible for the implementation of the training and validation of your accuracy predictor. The training process will take roughly 1-2 minutes.

Hint: you may refer to Tutorial 2 on how to train a neural network with PyTorch.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39


criterion = torch.nn.L1Loss().to(device)
optimizer = torch.optim.Adam(acc_predictor.parameters())
# the default value is zero
acc_predictor.base_acc.data += base_acc
for epoch in tqdm(range(10)):
 acc_predictor.train()
 for (data, label) in tqdm(train_loader, desc="Epoch%d" % (epoch + 1), position=0, leave=True):
 # step 1. Move the data and labels to device (cuda:0).
 data = data.to(device)
 label = label.to(device)
 ############### YOUR CODE STARTS HERE ###############
 # step 2. Run forward pass.
 pred = acc_predictor(data)
 # step 3. Calculate the loss.
 loss = criterion(pred, label)
 # step 4. Perform the backward pass.
 optimizer.zero_grad()
 loss.backward()
 optimizer.step()
 ################ YOUR CODE ENDS HERE ################

 acc_predictor.eval()
 with torch.no_grad():
 with tqdm(total=len(valid_loader), desc="Val", position=0, leave=True) as t:
 for (data, label) in valid_loader:
 # step 1. Move the data and labels to device (cuda:0).
 data = data.to(device)
 label = label.to(device)
 ############### YOUR CODE STARTS HERE ###############
 # step 2. Run forward pass.
 pred = acc_predictor(data)
 # step 3. Calculate the loss.
 loss = criterion(pred, label)
 ############### YOUR CODE ENDS HERE ###############
 t.set_postfix({"loss": loss.item()})
 t.update(1)

if not os.path.exists(acc_pred_checkpoint_path):
 torch.save(acc_predictor.cpu().state_dict(), acc_pred_checkpoint_path)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21


Epoch1: 100%|██████████| 157/157 [00:01<00:00, 101.01it/s]
Val: 100%|██████████| 40/40 [00:00<00:00, 48.38it/s, loss=0.00264]
Epoch2: 100%|██████████| 157/157 [00:01<00:00, 117.19it/s]
Val: 100%|██████████| 40/40 [00:00<00:00, 48.89it/s, loss=0.00161]
Epoch3: 100%|██████████| 157/157 [00:01<00:00, 115.86it/s]
Val: 100%|██████████| 40/40 [00:00<00:00, 49.96it/s, loss=0.00226]
Epoch4: 100%|██████████| 157/157 [00:02<00:00, 77.42it/s] 
Val: 100%|██████████| 40/40 [00:01<00:00, 37.70it/s, loss=0.00163]
Epoch5: 100%|██████████| 157/157 [00:01<00:00, 119.16it/s]
Val: 100%|██████████| 40/40 [00:00<00:00, 49.85it/s, loss=0.00174]
Epoch6: 100%|██████████| 157/157 [00:01<00:00, 122.14it/s]
Val: 100%|██████████| 40/40 [00:00<00:00, 49.41it/s, loss=0.00186]
Epoch7: 100%|██████████| 157/157 [00:01<00:00, 117.80it/s]
Val: 100%|██████████| 40/40 [00:00<00:00, 49.27it/s, loss=0.00148]
Epoch8: 100%|██████████| 157/157 [00:01<00:00, 117.19it/s]
Val: 100%|██████████| 40/40 [00:00<00:00, 49.47it/s, loss=0.0015]
Epoch9: 100%|██████████| 157/157 [00:01<00:00, 111.50it/s]
Val: 100%|██████████| 40/40 [00:01<00:00, 30.93it/s, loss=0.00132]
Epoch10: 100%|██████████| 157/157 [00:01<00:00, 89.64it/s] 
Val: 100%|██████████| 40/40 [00:00<00:00, 48.22it/s, loss=0.00143]
100%|██████████| 10/10 [00:23<00:00, 2.37s/it]

Now let’s plot the correlation of predicted accuracy against ground truth accuracy and make sure our predictor is reliable. To receive full score, you are expected to see a linear correlation in this part.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21


predicted_accuracies = []
ground_truth_accuracies = []
acc_predictor = acc_predictor.to("cuda:0")
acc_predictor.eval()
with torch.no_grad():
 with tqdm(total=len(valid_loader), desc="Val") as t:
 for (data, label) in valid_loader:
 data = data.to(device)
 label = label.to(device)
 pred = acc_predictor(data)
 predicted_accuracies += pred.cpu().numpy().tolist()
 ground_truth_accuracies += label.cpu().numpy().tolist()
 if len(predicted_accuracies) > 200:
 break
plt.scatter(predicted_accuracies, ground_truth_accuracies)
# draw y = x
min_acc, max_acc = min(predicted_accuracies), max(predicted_accuracies)
plt.plot([min_acc, max_acc], [min_acc, max_acc], c="red", linewidth=2)
plt.xlabel("Predicted accuracy")
plt.ylabel("Measured accuracy")
plt.title("Correlation between predicted accuracy and real accuracy")

1

Val: 0%| | 0/40 [00:00<?, ?it/s]

1

Text(0.5, 1.0, 'Correlation between predicted accuracy and real accuracy')

1

<Figure size 640x480 with 1 Axes>

Part 2. Neural Architecture Search (6 Questions, 65 pts + 10 bonus pts)

So far, we have defined both the efficiency and accuracy predictors. Let’s start fast model specialization with these two powerful predictors!

In this part, you are required to implement two typical search algorithms: random search and evolutionary search. The search algorithm aims to find the model architecture that provides the best accuracy while satisfying the efficiency constraints (e.g., MACs, peak memory).

Question 5 (5 pts): Complete the following random search agent.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28


class RandomSearcher:
 def __init__(self, efficiency_predictor, accuracy_predictor):
 self.efficiency_predictor = efficiency_predictor
 self.accuracy_predictor = accuracy_predictor

 def random_valid_sample(self, constraint):
 # randomly sample subnets until finding one that satisfies the constraint
 while True:
 sample = self.accuracy_predictor.arch_encoder.random_sample_arch()
 efficiency = self.efficiency_predictor.get_efficiency(sample)
 if self.efficiency_predictor.satisfy_constraint(efficiency, constraint):
 return sample, efficiency

 def run_search(self, constraint, n_subnets=100):
 subnet_pool = []
 # sample subnets
 for _ in tqdm(range(n_subnets)):
 sample, efficiency = self.random_valid_sample(constraint)
 subnet_pool.append(sample)
 # predict the accuracy of subnets
 accs = self.accuracy_predictor.predict_acc(subnet_pool)
 ############### YOUR CODE STARTS HERE ###############
 # hint: one line of code
 # get the index of the best subnet
 best_idx = accs.argmax()
 ############### YOUR CODE ENDS HERE #################
 # return the best subnet
 return accs[best_idx], subnet_pool[best_idx]

Question 6 (5 pts): Complete the following function.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19


def search_and_measure_acc(agent, constraint, **kwargs):
 ############### YOUR CODE STARTS HERE ###############
 # hint: call the search function
 best_info = agent.run_search(constraint=constraint, **kwargs)
 ############### YOUR CODE ENDS HERE #################
 # get searched subnet
 ofa_network.set_active_subnet(**best_info[1])
 subnet = ofa_network.get_active_subnet().to(device)
 # calibrate bn
 calib_bn(subnet, data_dir, 128, best_info[1]["image_size"])
 # build val loader
 val_loader = build_val_data_loader(data_dir, best_info[1]["image_size"], 128)
 # measure accuracy
 acc = validate(subnet, val_loader)
 # print best_info
 print(f"Accuracy of the selected subnet: {acc}")
 # visualize model architecture
 visualize_subnet(best_info[1])
 return acc, subnet

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16


random.seed(1)
np.random.seed(1)
nas_agent = RandomSearcher(efficiency_predictor, acc_predictor)
# MACs-constrained search
subnets_rs_macs = {}
for millonMACs in [50, 100]:
 search_constraint = dict(millonMACs=millonMACs)
 print(f"Random search with constraint: MACs <= {millonMACs}M")
 subnets_rs_macs[millonMACs] = search_and_measure_acc(nas_agent, search_constraint, n_subnets=300)

# memory-constrained search
subnets_rs_memory = {}
for KBPeakMemory in [256, 512]:
 search_constraint = dict(KBPeakMemory=KBPeakMemory)
 print(f"Random search with constraint: Peak memory <= {KBPeakMemory}KB")
 subnets_rs_memory[KBPeakMemory] = search_and_measure_acc(nas_agent, search_constraint, n_subnets=300)

1

Random search with constraint: MACs <= 50M

1
2


100%|██████████| 300/300 [01:06<00:00, 4.49it/s]
Validate: 100%|██████████| 32/32 [00:08<00:00, 3.60it/s, loss=0.187, top1=93.3]

1

Accuracy of the selected subnet: 93.27543427346657

1

<Figure size 1630.4x177.2 with 1 Axes>

1

Random search with constraint: MACs <= 100M

1
2


100%|██████████| 300/300 [01:06<00:00, 4.50it/s]
Validate: 100%|██████████| 32/32 [00:09<00:00, 3.46it/s, loss=0.188, top1=93.5]

1

Accuracy of the selected subnet: 93.5483871270646

1

<Figure size 1630.4x177.2 with 1 Axes>

1

Random search with constraint: Peak memory <= 256KB

1
2


100%|██████████| 300/300 [02:24<00:00, 2.08it/s]
Validate: 100%|██████████| 32/32 [00:10<00:00, 3.17it/s, loss=0.206, top1=92.8]

1

Accuracy of the selected subnet: 92.77915630482562

1

<Figure size 1561.2x177.2 with 1 Axes>

1

Random search with constraint: Peak memory <= 512KB

1
2


100%|██████████| 300/300 [01:12<00:00, 4.11it/s]
Validate: 100%|██████████| 32/32 [00:10<00:00, 3.11it/s, loss=0.19, top1=93.2]

1

Accuracy of the selected subnet: 93.15136479455839

1

<Figure size 1561.2x177.2 with 1 Axes>

Question7 (20 pts): Complete the following evolutionary search agent.

Now you have succesfully implemented the random search algorithm. In this part, we will implement a more sample-efficient search algorithm, evolutionary search. Evolutionary search is inspired by the evolution algorithm (or genetic algorithm). A population of sub-networks are first sampled from the design space. Then, in each generation, we perform random mutation and crossover operations as is shown in the figure above. The sub-networks with highest accuracy will be kept, and this process will be repeated until the number of generations reaches max_time_budget. Similar to the random search, throughout the search process, all sub-networks that cannot satisfy the efficiency constraint will be discarded.

 1
 2
 3
 4
 5
 6
 7
 8
 9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121


class EvolutionSearcher:
 def __init__(self, efficiency_predictor, accuracy_predictor, **kwargs):
 self.efficiency_predictor = efficiency_predictor
 self.accuracy_predictor = accuracy_predictor

 # evolution hyper-parameters
 self.arch_mutate_prob = kwargs.get("arch_mutate_prob", 0.1)
 self.resolution_mutate_prob = kwargs.get("resolution_mutate_prob", 0.5)
 self.population_size = kwargs.get("population_size", 100)
 self.max_time_budget = kwargs.get("max_time_budget", 500)
 self.parent_ratio = kwargs.get("parent_ratio", 0.25)
 self.mutation_ratio = kwargs.get("mutation_ratio", 0.5)

 def update_hyper_params(self, new_param_dict):
 self.__dict__.update(new_param_dict)

 def random_valid_sample(self, constraint):
 # randomly sample subnets until finding one that satisfies the constraint
 while True:
 sample = self.accuracy_predictor.arch_encoder.random_sample_arch()
 efficiency = self.efficiency_predictor.get_efficiency(sample)
 if self.efficiency_predictor.satisfy_constraint(efficiency, constraint):
 return sample, efficiency

 def mutate_sample(self, sample, constraint):
 while True:
 new_sample = copy.deepcopy(sample)

 self.accuracy_predictor.arch_encoder.mutate_resolution(new_sample, self.resolution_mutate_prob)
 self.accuracy_predictor.arch_encoder.mutate_width(new_sample, self.arch_mutate_prob)
 self.accuracy_predictor.arch_encoder.mutate_arch(new_sample, self.arch_mutate_prob)

 efficiency = self.efficiency_predictor.get_efficiency(new_sample)
 if self.efficiency_predictor.satisfy_constraint(efficiency, constraint):
 return new_sample, efficiency

 def crossover_sample(self, sample1, sample2, constraint):
 while True:
 new_sample = copy.deepcopy(sample1)
 for key in new_sample.keys():
 if not isinstance(new_sample[key], list):
 ############### YOUR CODE STARTS HERE ###############
 # hint: randomly choose the value from sample1[key] and sample2[key], random.choice
 new_sample[key] = random.choice([sample1[key], sample2[key]])
 ############### YOUR CODE ENDS HERE #################
 else:
 for i in range(len(new_sample[key])):
 ############### YOUR CODE STARTS HERE ###############
 new_sample[key][i] = random.choice([sample1[key][i], sample2[key][i]])
 ############### YOUR CODE ENDS HERE #################

 efficiency = self.efficiency_predictor.get_efficiency(new_sample)
 if self.efficiency_predictor.satisfy_constraint(efficiency, constraint):
 return new_sample, efficiency

 def run_search(self, constraint, **kwargs):
 self.update_hyper_params(kwargs)

 mutation_numbers = int(round(self.mutation_ratio * self.population_size))
 parents_size = int(round(self.parent_ratio * self.population_size))

 best_valids = [-100]
 population = [] # (acc, sample) tuples
 child_pool = []
 best_info = None
 # generate random population
 for _ in range(self.population_size):
 sample, efficiency = self.random_valid_sample(constraint)
 child_pool.append(sample)

 accs = self.accuracy_predictor.predict_acc(child_pool)
 for i in range(self.population_size):
 population.append((accs[i].item(), child_pool[i]))

 # evolving the population
 with tqdm(total=self.max_time_budget) as t:
 for i in range(self.max_time_budget):
 ############### YOUR CODE STARTS HERE ###############
 # hint: sort the population according to the acc (descending order)
 population = sorted(population, key=lambda x: x[0], reverse=True)
 ############### YOUR CODE ENDS HERE #################

 ############### YOUR CODE STARTS HERE ###############
 # hint: keep topK samples in the population, K = parents_size
 # the others are discarded.
 population = population[:parents_size]
 ############### YOUR CODE ENDS HERE #################

 # update best info
 acc = population[0][0]
 if acc > best_valids[-1]:
 best_valids.append(acc)
 best_info = population[0]
 else:
 best_valids.append(best_valids[-1])

 child_pool = []
 for j in range(mutation_numbers):
 # randomly choose a sample
 par_sample = population[np.random.randint(parents_size)][1]
 # mutate this sample
 new_sample, efficiency = self.mutate_sample(par_sample, constraint)
 child_pool.append(new_sample)

 for j in range(self.population_size - mutation_numbers):
 # randomly choose two samples
 par_sample1 = population[np.random.randint(parents_size)][1]
 par_sample2 = population[np.random.randint(parents_size)][1]
 # crossover
 new_sample, efficiency = self.crossover_sample(
 par_sample1, par_sample2, constraint
 )
 child_pool.append(new_sample)
 # predict accuracy with the accuracy predictor
 accs = self.accuracy_predictor.predict_acc(child_pool)
 for j in range(self.population_size):
 population.append((accs[j].item(), child_pool[j]))

 t.update(1)

 return best_info

Question 8 (10pts): Run evolutionary search and tune evo_params to optimize the results. Describe your findings.

Answer:

The default population size and time budget are too small. Increasing them (e.g., population_size=50, max_time_budget=20) improves results but increases search cost (~4 min per search in our run).
Increasing resolution_mutate_prob to 0.5 helps, consistent with Question 1. With tuned params, evolutionary search reached ~92.3–93.6% validation accuracy under MAC/memory constraints (e.g., 92.3% at ≤50M MACs, 93.6% at ≤512KB peak memory).

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27


random.seed(1)
np.random.seed(1)

# hint: tune hyper-parameters below
evo_params = {
 'arch_mutate_prob': 0.1, # The probability of architecture mutation in evolutionary search
 'resolution_mutate_prob': 0.5, # The probability of resolution mutation in evolutionary search
 'population_size': 50, # The size of the population
 'max_time_budget': 20,
 'parent_ratio': 0.25,
 'mutation_ratio': 0.3,
}

nas_agent = EvolutionSearcher(efficiency_predictor, acc_predictor, **evo_params)
# MACs-constrained search
subnets_evo_macs = {}
for millonMACs in [50, 100]:
 search_constraint = dict(millionMACs=millonMACs)
 print(f"Evolutionary search with constraint: MACs <= {millonMACs}M")
 subnets_evo_macs[millonMACs] = search_and_measure_acc(nas_agent, search_constraint)

# memory-constrained search
subnets_evo_memory = {}
for KBPeakMemory in [256, 512]:
 search_constraint = dict(KBPeakMemory=KBPeakMemory)
 print(f"Evolutionary search with constraint: Peak memory <= {KBPeakMemory}KB")
 subnets_evo_memory[KBPeakMemory] = search_and_measure_acc(nas_agent, search_constraint)

1

Evolutionary search with constraint: MACs <= 50M

1
2


100%|██████████| 20/20 [04:13<00:00, 12.66s/it]
Validate: 100%|██████████| 32/32 [00:09<00:00, 3.54it/s, loss=0.209, top1=92.3]

1

Accuracy of the selected subnet: 92.25806454641943

1

<Figure size 1422.4x177.2 with 1 Axes>

1

Evolutionary search with constraint: MACs <= 100M

1
2


100%|██████████| 20/20 [04:11<00:00, 12.57s/it]
Validate: 100%|██████████| 32/32 [00:10<00:00, 3.18it/s, loss=0.19, top1=92.9]

1

Accuracy of the selected subnet: 92.92803973252366

1

<Figure size 1422.4x177.2 with 1 Axes>

1

Evolutionary search with constraint: Peak memory <= 256KB

1
2


100%|██████████| 20/20 [04:21<00:00, 13.08s/it]
Validate: 100%|██████████| 32/32 [00:08<00:00, 3.62it/s, loss=0.196, top1=92.8]

1

Accuracy of the selected subnet: 92.75434240904399

1

<Figure size 1699.6x177.2 with 1 Axes>

1

Evolutionary search with constraint: Peak memory <= 512KB

1
2


100%|██████████| 20/20 [04:12<00:00, 12.60s/it]
Validate: 100%|██████████| 32/32 [00:09<00:00, 3.43it/s, loss=0.187, top1=93.6]

1

Accuracy of the selected subnet: 93.6228288144095

1

<Figure size 1630.4x177.2 with 1 Axes>

Question 9 (15 pts + 10 pts bonus): Run evolutionary search under real-world constraints.

In real-world applications, we may have multiple efficiency constraints: https://blog.tensorflow.org/2019/10/visual-wake-words-with-tensorflow-lite_30.html. Use evolutionary search to find models that satisfy the following constraints:

[15 pts] 250 KB, 60M MACs (acc >= 92.5% to get the full credit)
[10 pts, bonus] 200KB, 30M MACs (acc >= 90% to get the full credit)

Hint: You do not have to use the same evo_params for these two tasks.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18


random.seed(1)
np.random.seed(1)
# hint: tune hyper-parameters below
evo_params = {
 'arch_mutate_prob': 0.1, # The probability of architecture mutation in evolutionary search
 'resolution_mutate_prob': 0.5, # The probability of resolution mutation in evolutionary search
 'population_size': 50, # The size of the population
 'max_time_budget': 20,
 'parent_ratio': 0.25,
 'mutation_ratio': 0.3,
}

nas_agent = EvolutionSearcher(efficiency_predictor, acc_predictor, **evo_params)

(millionMACs, KBPeakMemory) = [60, 250]
print(f"Evolution search with constraint: MACs <= {millionMACs}M, peak memory <= {KBPeakMemory}KB")
search_and_measure_acc(nas_agent, dict(millionMACs=millionMACs, KBPeakMemory=KBPeakMemory))
print("Evolution search finished!")

1

Evolution search with constraint: MACs <= 60M, peak memory <= 250KB

1
2


100%|██████████| 20/20 [04:32<00:00, 13.65s/it]
Validate: 100%|██████████| 32/32 [00:09<00:00, 3.24it/s, loss=0.203, top1=92.7]

1

Accuracy of the selected subnet: 92.72952851326235

1

<Figure size 1561.2x177.2 with 1 Axes>

1

Evolution search finished!

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18


random.seed(1)
np.random.seed(1)
# hint: tune hyper-parameters below
evo_params = {
 'arch_mutate_prob': 0.1, # The probability of architecture mutation in evolutionary search
 'resolution_mutate_prob': 0.5, # The probability of resolution mutation in evolutionary search
 'population_size': 50, # The size of the population
 'max_time_budget': 30,
 'parent_ratio': 0.25,
 'mutation_ratio': 0.3,
}

nas_agent = EvolutionSearcher(efficiency_predictor, acc_predictor, **evo_params)

(millionMACs, KBPeakMemory) = [30, 200]
print(f"Evolution search with constraint: MACs <= {millionMACs}M, peak memory <= {KBPeakMemory}KB")
search_and_measure_acc(nas_agent, dict(millionMACs=millionMACs, KBPeakMemory=KBPeakMemory))
print("Evolution search finished!")

1

Evolution search with constraint: MACs <= 30M, peak memory <= 200KB

1
2


100%|██████████| 30/30 [05:51<00:00, 11.71s/it]
Validate: 100%|██████████| 32/32 [00:07<00:00, 4.02it/s, loss=0.248, top1=90.5]

1

Accuracy of the selected subnet: 90.47146403878261

1

<Figure size 1353.2x177.2 with 1 Axes>

1

Evolution search finished!

Question 10 (10 pts): Is it possible to find a subnet with the following efficiency constraints in the current design space?

A: The activation size of the subnet is at most 256KB and the MACs of the subnet is at most 15M.
B: The activation size of the subnet is at most 64 KB.

Answer:

A (≤256 KB activation, ≤15M MACs): Yes. The smallest subnet in the design space uses ~8.3M MACs and ~72 KB peak activation (within both limits) and achieves ~83.4% accuracy at 96×96.
B (≤64 KB activation): No. Even the smallest subnet has ~72 KB peak activation, which already exceeds the 64 KB limit, so no feasible subnet exists in this design space.

MIT 6.5940 EfficientML.ai Lab 2: Quantization

Tue, 19 May 2026 00:00:00 +0000

MIT 6.5940 EfficientML.ai Lab 2: Quantization

This colab notebook provides code and a framework for Lab 2 quantization. You can work out your solutions here.

Please fill out this feedback form when you finished this lab. We would love to hear your thoughts or feedback on how we can improve this lab!

Goals

In this assignment, you will practice quantizing a classical neural network model to reduce both model size and latency. The goals of this assignment are as follows:

Understand the basic concept of quantization
Implement and apply k-means quantization
Implement and apply quantization-aware training for k-means quantization
Implement and apply linear quantization
Implement and apply integer-only inference for linear quantization
Get a basic understanding of performance improvement (such as speedup) from quantization
Understand the differences and tradeoffs between these quantization approaches

There are 2 main sections: K-Means Quantization and Linear Quantization.

There are 10 questions in total:

For K-Means Quantization, there are 3 questions (Question 1-3).
For Linear Quantization, there are 6 questions (Question 4-9).
Question 10 compares k-means quantization and linear quantization.

Setup

First, install the required packages and download the datasets and pretrained model. Here we use CIFAR10 dataset and VGG network which is the same as what we used in the Lab 0 tutorial.

1
2
3
4
5


print('Installing torchprofile...')
!pip install torchprofile 1>/dev/null
print('Installing fast-pytorch-kmeans...')
! pip install fast-pytorch-kmeans 1>/dev/null
print('All required packages have been successfully installed!')

1
2
3


Installing torchprofile...
Installing fast-pytorch-kmeans...
All required packages have been successfully installed!

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24


import copy
import math
import random
from collections import OrderedDict, defaultdict

from matplotlib import pyplot as plt
from matplotlib.colors import ListedColormap
import numpy as np
from tqdm.auto import tqdm

import torch
from torch import nn
from torch.optim import *
from torch.optim.lr_scheduler import *
from torch.utils.data import DataLoader
from torchprofile import profile_macs
from torchvision.datasets import *
from torchvision.transforms import *

from torchprofile import profile_macs

assert torch.cuda.is_available(), \
"The current runtime does not have CUDA support." \
"Please go to menu bar (Runtime - Change runtime type) and select GPU"

1
2
3


random.seed(0)
np.random.seed(0)
torch.manual_seed(0)

1

<torch._C.Generator at 0x7b68d1410a90>

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19


def download_url(url, model_dir='.', overwrite=False):
 import os, sys
 from urllib.request import urlretrieve
 target_dir = url.split('/')[-1]
 model_dir = os.path.expanduser(model_dir)
 try:
 if not os.path.exists(model_dir):
 os.makedirs(model_dir)
 model_dir = os.path.join(model_dir, target_dir)
 cached_file = model_dir
 if not os.path.exists(cached_file) or overwrite:
 sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file))
 urlretrieve(url, cached_file)
 return cached_file
 except Exception as e:
 # remove lock file so download can be executed next time.
 os.remove(os.path.join(model_dir, 'download.lock'))
 sys.stderr.write('Failed to download from url %s' % url + '\n' + str(e) + '\n')
 return None

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39


class VGG(nn.Module):
 ARCH = [64, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M']

 def __init__(self) -> None:
 super().__init__()

 layers = []
 counts = defaultdict(int)

 def add(name: str, layer: nn.Module) -> None:
 layers.append((f"{name}{counts[name]}", layer))
 counts[name] += 1

 in_channels = 3
 for x in self.ARCH:
 if x != 'M':
 # conv-bn-relu
 add("conv", nn.Conv2d(in_channels, x, 3, padding=1, bias=False))
 add("bn", nn.BatchNorm2d(x))
 add("relu", nn.ReLU(True))
 in_channels = x
 else:
 # maxpool
 add("pool", nn.MaxPool2d(2))
 add("avgpool", nn.AvgPool2d(2))
 self.backbone = nn.Sequential(OrderedDict(layers))
 self.classifier = nn.Linear(512, 10)

 def forward(self, x: torch.Tensor) -> torch.Tensor:
 # backbone: [N, 3, 32, 32] => [N, 512, 2, 2]
 x = self.backbone(x)

 # avgpool: [N, 512, 2, 2] => [N, 512]
 # x = x.mean([2, 3])
 x = x.view(x.shape[0], -1)

 # classifier: [N, 512] => [N, 10]
 x = self.classifier(x)
 return x

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32


def train(
 model: nn.Module,
 dataloader: DataLoader,
 criterion: nn.Module,
 optimizer: Optimizer,
 scheduler: LambdaLR,
 callbacks = None
) -> None:
 model.train()

 for inputs, targets in tqdm(dataloader, desc='train', leave=False):
 # Move the data from CPU to GPU
 inputs = inputs.cuda()
 targets = targets.cuda()

 # Reset the gradients (from the last iteration)
 optimizer.zero_grad()

 # Forward inference
 outputs = model(inputs)
 loss = criterion(outputs, targets)

 # Backward propagation
 loss.backward()

 # Update optimizer and LR scheduler
 optimizer.step()
 scheduler.step()

 if callbacks is not None:
 for callback in callbacks:
 callback()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31


@torch.inference_mode()
def evaluate(
 model: nn.Module,
 dataloader: DataLoader,
 extra_preprocess = None
) -> float:
 model.eval()

 num_samples = 0
 num_correct = 0

 for inputs, targets in tqdm(dataloader, desc="eval", leave=False):
 # Move the data from CPU to GPU
 inputs = inputs.cuda()
 if extra_preprocess is not None:
 for preprocess in extra_preprocess:
 inputs = preprocess(inputs)

 targets = targets.cuda()

 # Inference
 outputs = model(inputs)

 # Convert logits to class indices
 outputs = outputs.argmax(dim=1)

 # Update metrics
 num_samples += targets.size(0)
 num_correct += (outputs == targets).sum()

 return (num_correct / num_samples * 100).item()

Helpler Functions (Flops, Model Size calculation, etc.)

1
2
3


def get_model_flops(model, inputs):
 num_macs = profile_macs(model, inputs)
 return num_macs

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14


def get_model_size(model: nn.Module, data_width=32):
 """
 calculate the model size in bits
 :param data_width: #bits per element
 """
 num_elements = 0
 for param in model.parameters():
 num_elements += param.numel()
 return num_elements * data_width

Byte = 8
KiB = 1024 * Byte
MiB = 1024 * KiB
GiB = 1024 * MiB

Define misc funcions for verification.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37


def test_k_means_quantize(
 test_tensor=torch.tensor([
 [-0.3747, 0.0874, 0.3200, -0.4868, 0.4404],
 [-0.0402, 0.2322, -0.2024, -0.4986, 0.1814],
 [ 0.3102, -0.3942, -0.2030, 0.0883, -0.4741],
 [-0.1592, -0.0777, -0.3946, -0.2128, 0.2675],
 [ 0.0611, -0.1933, -0.4350, 0.2928, -0.1087]]),
 bitwidth=2):
 def plot_matrix(tensor, ax, title, cmap=ListedColormap(['white'])):
 ax.imshow(tensor.cpu().numpy(), vmin=-0.5, vmax=0.5, cmap=cmap)
 ax.set_title(title)
 ax.set_yticklabels([])
 ax.set_xticklabels([])
 for i in range(tensor.shape[1]):
 for j in range(tensor.shape[0]):
 text = ax.text(j, i, f'{tensor[i, j].item():.2f}',
 ha="center", va="center", color="k")

 fig, axes = plt.subplots(1,2, figsize=(8, 12))
 ax_left, ax_right = axes.ravel()

 print(test_tensor)
 plot_matrix(test_tensor, ax_left, 'original tensor')

 num_unique_values_before_quantization = test_tensor.unique().numel()
 k_means_quantize(test_tensor, bitwidth=bitwidth)
 num_unique_values_after_quantization = test_tensor.unique().numel()
 print('* Test k_means_quantize()')
 print(f' target bitwidth: {bitwidth} bits')
 print(f' num unique values before k-means quantization: {num_unique_values_before_quantization}')
 print(f' num unique values after k-means quantization: {num_unique_values_after_quantization}')
 assert num_unique_values_after_quantization == min((1 << bitwidth), num_unique_values_before_quantization)
 print('* Test passed.')

 plot_matrix(test_tensor, ax_right, f'{bitwidth}-bit k-means quantized tensor', cmap='tab20c')
 fig.tight_layout()
 plt.show()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46


def test_linear_quantize(
 test_tensor=torch.tensor([
 [ 0.0523, 0.6364, -0.0968, -0.0020, 0.1940],
 [ 0.7500, 0.5507, 0.6188, -0.1734, 0.4677],
 [-0.0669, 0.3836, 0.4297, 0.6267, -0.0695],
 [ 0.1536, -0.0038, 0.6075, 0.6817, 0.0601],
 [ 0.6446, -0.2500, 0.5376, -0.2226, 0.2333]]),
 quantized_test_tensor=torch.tensor([
 [-1, 1, -1, -1, 0],
 [ 1, 1, 1, -2, 0],
 [-1, 0, 0, 1, -1],
 [-1, -1, 1, 1, -1],
 [ 1, -2, 1, -2, 0]], dtype=torch.int8),
 real_min=-0.25, real_max=0.75, bitwidth=2, scale=1/3, zero_point=-1):
 def plot_matrix(tensor, ax, title, vmin=0, vmax=1, cmap=ListedColormap(['white'])):
 ax.imshow(tensor.cpu().numpy(), vmin=vmin, vmax=vmax, cmap=cmap)
 ax.set_title(title)
 ax.set_yticklabels([])
 ax.set_xticklabels([])
 for i in range(tensor.shape[0]):
 for j in range(tensor.shape[1]):
 datum = tensor[i, j].item()
 if isinstance(datum, float):
 text = ax.text(j, i, f'{datum:.2f}',
 ha="center", va="center", color="k")
 else:
 text = ax.text(j, i, f'{datum}',
 ha="center", va="center", color="k")
 quantized_min, quantized_max = get_quantized_range(bitwidth)
 fig, axes = plt.subplots(1,3, figsize=(10, 32))
 plot_matrix(test_tensor, axes[0], 'original tensor', vmin=real_min, vmax=real_max)
 _quantized_test_tensor = linear_quantize(
 test_tensor, bitwidth=bitwidth, scale=scale, zero_point=zero_point)
 _reconstructed_test_tensor = scale * (_quantized_test_tensor.float() - zero_point)
 print('* Test linear_quantize()')
 print(f' target bitwidth: {bitwidth} bits')
 print(f' scale: {scale}')
 print(f' zero point: {zero_point}')
 assert _quantized_test_tensor.equal(quantized_test_tensor)
 print('* Test passed.')
 plot_matrix(_quantized_test_tensor, axes[1], f'2-bit linear quantized tensor',
 vmin=quantized_min, vmax=quantized_max, cmap='tab20c')
 plot_matrix(_reconstructed_test_tensor, axes[2], f'reconstructed tensor',
 vmin=real_min, vmax=real_max, cmap='tab20c')
 fig.tight_layout()
 plt.show()

 1
 2
 3
 4
 5
 6
 7
 8
 9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100


def test_quantized_fc(
 input=torch.tensor([
 [0.6118, 0.7288, 0.8511, 0.2849, 0.8427, 0.7435, 0.4014, 0.2794],
 [0.3676, 0.2426, 0.1612, 0.7684, 0.6038, 0.0400, 0.2240, 0.4237],
 [0.6565, 0.6878, 0.4670, 0.3470, 0.2281, 0.8074, 0.0178, 0.3999],
 [0.1863, 0.3567, 0.6104, 0.0497, 0.0577, 0.2990, 0.6687, 0.8626]]),
 weight=torch.tensor([
 [ 1.2626e-01, -1.4752e-01, 8.1910e-02, 2.4982e-01, -1.0495e-01,
 -1.9227e-01, -1.8550e-01, -1.5700e-01],
 [ 2.7624e-01, -4.3835e-01, 5.1010e-02, -1.2020e-01, -2.0344e-01,
 1.0202e-01, -2.0799e-01, 2.4112e-01],
 [-3.8216e-01, -2.8047e-01, 8.5238e-02, -4.2504e-01, -2.0952e-01,
 3.2018e-01, -3.3619e-01, 2.0219e-01],
 [ 8.9233e-02, -1.0124e-01, 1.1467e-01, 2.0091e-01, 1.1438e-01,
 -4.2427e-01, 1.0178e-01, -3.0941e-04],
 [-1.8837e-02, -2.1256e-01, -4.5285e-01, 2.0949e-01, -3.8684e-01,
 -1.7100e-01, -4.5331e-01, -2.0433e-01],
 [-2.0038e-01, -5.3757e-02, 1.8997e-01, -3.6866e-01, 5.5484e-02,
 1.5643e-01, -2.3538e-01, 2.1103e-01],
 [-2.6875e-01, 2.4984e-01, -2.3514e-01, 2.5527e-01, 2.0322e-01,
 3.7675e-01, 6.1563e-02, 1.7201e-01],
 [ 3.3541e-01, -3.3555e-01, -4.3349e-01, 4.3043e-01, -2.0498e-01,
 -1.8366e-01, -9.1553e-02, -4.1168e-01]]),
 bias=torch.tensor([ 0.1954, -0.2756, 0.3113, 0.1149, 0.4274, 0.2429, -0.1721, -0.2502]),
 quantized_bias=torch.tensor([ 3, -2, 3, 1, 3, 2, -2, -2], dtype=torch.int32),
 shifted_quantized_bias=torch.tensor([-1, 0, -3, -1, -3, 0, 2, -4], dtype=torch.int32),
 calc_quantized_output=torch.tensor([
 [ 0, -1, 0, -1, -1, 0, 1, -2],
 [ 0, 0, -1, 0, 0, 0, 0, -1],
 [ 0, 0, 0, -1, 0, 0, 0, -1],
 [ 0, 0, 0, 0, 0, 1, -1, -2]], dtype=torch.int8),
 bitwidth=2, batch_size=4, in_channels=8, out_channels=8):
 def plot_matrix(tensor, ax, title, vmin=0, vmax=1, cmap=ListedColormap(['white'])):
 ax.imshow(tensor.cpu().numpy(), vmin=vmin, vmax=vmax, cmap=cmap)
 ax.set_title(title)
 ax.set_yticklabels([])
 ax.set_xticklabels([])
 for i in range(tensor.shape[0]):
 for j in range(tensor.shape[1]):
 datum = tensor[i, j].item()
 if isinstance(datum, float):
 text = ax.text(j, i, f'{datum:.2f}',
 ha="center", va="center", color="k")
 else:
 text = ax.text(j, i, f'{datum}',
 ha="center", va="center", color="k")

 output = torch.nn.functional.linear(input, weight, bias)

 quantized_weight, weight_scale, weight_zero_point = \
 linear_quantize_weight_per_channel(weight, bitwidth)
 quantized_input, input_scale, input_zero_point = \
 linear_quantize_feature(input, bitwidth)
 _quantized_bias, bias_scale, bias_zero_point = \
 linear_quantize_bias_per_output_channel(bias, weight_scale, input_scale)
 assert _quantized_bias.equal(_quantized_bias)
 _shifted_quantized_bias = \
 shift_quantized_linear_bias(quantized_bias, quantized_weight, input_zero_point)
 assert _shifted_quantized_bias.equal(shifted_quantized_bias)
 quantized_output, output_scale, output_zero_point = \
 linear_quantize_feature(output, bitwidth)

 _calc_quantized_output = quantized_linear(
 quantized_input, quantized_weight, shifted_quantized_bias,
 bitwidth, bitwidth,
 input_zero_point, output_zero_point,
 input_scale, weight_scale, output_scale)
 assert _calc_quantized_output.equal(calc_quantized_output)

 reconstructed_weight = weight_scale * (quantized_weight.float() - weight_zero_point)
 reconstructed_input = input_scale * (quantized_input.float() - input_zero_point)
 reconstructed_bias = bias_scale * (quantized_bias.float() - bias_zero_point)
 reconstructed_calc_output = output_scale * (calc_quantized_output.float() - output_zero_point)

 fig, axes = plt.subplots(3,3, figsize=(15, 12))
 quantized_min, quantized_max = get_quantized_range(bitwidth)
 plot_matrix(weight, axes[0, 0], 'original weight', vmin=-0.5, vmax=0.5)
 plot_matrix(input.t(), axes[1, 0], 'original input', vmin=0, vmax=1)
 plot_matrix(output.t(), axes[2, 0], 'original output', vmin=-1.5, vmax=1.5)
 plot_matrix(quantized_weight, axes[0, 1], f'{bitwidth}-bit linear quantized weight',
 vmin=quantized_min, vmax=quantized_max, cmap='tab20c')
 plot_matrix(quantized_input.t(), axes[1, 1], f'{bitwidth}-bit linear quantized input',
 vmin=quantized_min, vmax=quantized_max, cmap='tab20c')
 plot_matrix(calc_quantized_output.t(), axes[2, 1], f'quantized output from quantized_linear()',
 vmin=quantized_min, vmax=quantized_max, cmap='tab20c')
 plot_matrix(reconstructed_weight, axes[0, 2], f'reconstructed weight',
 vmin=-0.5, vmax=0.5, cmap='tab20c')
 plot_matrix(reconstructed_input.t(), axes[1, 2], f'reconstructed input',
 vmin=0, vmax=1, cmap='tab20c')
 plot_matrix(reconstructed_calc_output.t(), axes[2, 2], f'reconstructed output',
 vmin=-1.5, vmax=1.5, cmap='tab20c')

 print('* Test quantized_fc()')
 print(f' target bitwidth: {bitwidth} bits')
 print(f' batch size: {batch_size}')
 print(f' input channels: {in_channels}')
 print(f' output channels: {out_channels}')
 print('* Test passed.')
 fig.tight_layout()
 plt.show()

Load Pretrained Model

1
2
3
4
5
6


checkpoint_url = "https://hanlab18.mit.edu/files/course/labs/vgg.cifar.pretrained.pth"
checkpoint = torch.load(download_url(checkpoint_url), map_location="cpu")
model = VGG().cuda()
print(f"=> loading checkpoint '{checkpoint_url}'")
model.load_state_dict(checkpoint['state_dict'])
recover_model = lambda : model.load_state_dict(checkpoint['state_dict'])

1

Downloading: "https://hanlab18.mit.edu/files/course/labs/vgg.cifar.pretrained.pth" to ./vgg.cifar.pretrained.pth

1

=> loading checkpoint 'https://hanlab18.mit.edu/files/course/labs/vgg.cifar.pretrained.pth'

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26


image_size = 32
transforms = {
 "train": Compose([
 RandomCrop(image_size, padding=4),
 RandomHorizontalFlip(),
 ToTensor(),
 ]),
 "test": ToTensor(),
}
dataset = {}
for split in ["train", "test"]:
 dataset[split] = CIFAR10(
 root="data/cifar10",
 train=(split == "train"),
 download=True,
 transform=transforms[split],
 )
dataloader = {}
for split in ['train', 'test']:
 dataloader[split] = DataLoader(
 dataset[split],
 batch_size=512,
 shuffle=(split == 'train'),
 num_workers=0,
 pin_memory=True,
 )

1

100%|██████████| 170M/170M [00:07<00:00, 22.2MB/s]

Let’s First Evaluate the Accuracy and Model Size of the FP32 Model

1
2
3
4


fp32_model_accuracy = evaluate(model, dataloader['test'])
fp32_model_size = get_model_size(model)
print(f"fp32 model has accuracy={fp32_model_accuracy:.2f}%")
print(f"fp32 model has size={fp32_model_size/MiB:.2f} MiB")

1

eval: 0%| | 0/20 [00:00<?, ?it/s]

1
2


fp32 model has accuracy=92.95%
fp32 model has size=35.20 MiB

K-Means Quantization

Network quantization compresses the network by reducing the bits per weight required to represent the deep network. The quantized network can have a faster inference speed with hardware support.

In this section, we will explore the K-means quantization for neural networks as in Deep Compression: Compressing Deep Neural Networks With Pruning, Trained Quantization And Huffman Coding.

A $n$-bit k-means quantization will divide synapses into $2^n$ clusters, and synapses in the same cluster will share the same weight value.

Therefore, k-means quantization will create a codebook, inlcuding

centroids: $2^n$ fp32 cluster centers.
labels: a $n$-bit integer tensor with the same #elements of the original fp32 weights tensor. Each integer indicates which cluster it belongs to.

During the inference, a fp32 tensor is generated based on the codebook for inference:

quantized_weight = codebook.centroids\[*codebook.labels*\].view_as(weight)

1
2
3


from collections import namedtuple

Codebook = namedtuple('Codebook', ['centroids', 'labels'])

Question 1 (10 pts)

Please complete the following K-Means quantization function.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31


from fast_pytorch_kmeans import KMeans

def k_means_quantize(fp32_tensor: torch.Tensor, bitwidth=4, codebook=None):
 """
 quantize tensor using k-means clustering
 :param fp32_tensor:
 :param bitwidth: [int] quantization bit width, default=4
 :param codebook: [Codebook] (the cluster centroids, the cluster label tensor)
 :return:
 [Codebook = (centroids, labels)]
 centroids: [torch.(cuda.)FloatTensor] the cluster centroids
 labels: [torch.(cuda.)LongTensor] cluster label tensor
 """
 if codebook is None:
 ############### YOUR CODE STARTS HERE ###############
 # get number of clusters based on the quantization precision
 # hint: one line of code
 n_clusters = 1 << bitwidth
 ############### YOUR CODE ENDS HERE #################
 # use k-means to get the quantization centroids
 kmeans = KMeans(n_clusters=n_clusters, mode='euclidean', verbose=0)
 labels = kmeans.fit_predict(fp32_tensor.view(-1, 1)).to(torch.long)
 centroids = kmeans.centroids.to(torch.float).view(-1)
 codebook = Codebook(centroids, labels)
 ############### YOUR CODE STARTS HERE ###############
 # decode the codebook into k-means quantized tensor for inference
 # hint: one line of code
 quantized_tensor = codebook.centroids[codebook.labels]
 ############### YOUR CODE ENDS HERE #################
 fp32_tensor.set_(quantized_tensor.view_as(fp32_tensor))
 return codebook

Let’s verify the functionality of defined k-means quantization by applying the function above on a dummy tensor.

1

test_k_means_quantize()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


tensor([[-0.3747, 0.0874, 0.3200, -0.4868, 0.4404],
 [-0.0402, 0.2322, -0.2024, -0.4986, 0.1814],
 [ 0.3102, -0.3942, -0.2030, 0.0883, -0.4741],
 [-0.1592, -0.0777, -0.3946, -0.2128, 0.2675],
 [ 0.0611, -0.1933, -0.4350, 0.2928, -0.1087]])
* Test k_means_quantize()
 target bitwidth: 2 bits
 num unique values before k-means quantization: 25
 num unique values after k-means quantization: 4
* Test passed.

1

<Figure size 800x1200 with 2 Axes>

Question 2 (10 pts)

The last code cell performs 2-bit k-means quantization and plots the tensor before and after the quantization. Each cluster is rendered with a unique color. There are 4 unique colors rendered in the quantized tensor.

Given this observation, please answer the following questions.

Question 2.1 (5 pts)

If 4-bit k-means quantization is performed, how many unique colors will be rendered in the quantized tensor?

16 (2^4 = 16 unique values/clusters for 4-bit quantization)

Question 2.2 (5 pts)

If n-bit k-means quantization is performed, how many unique colors will be rendered in the quantized tensor?

$2^n$ unique colors, where $n$ is the number of bits. For example, 2-bit has 4 colors ($2^2=4$), 4-bit has 16 colors ($2^4=16$), etc.

K-Means Quantization on Whole Model

Similar to what we did in lab 1, we now wrap the k-means quantization function into a class for quantizing the whole model. In class KMeansQuantizer, we have to keep a record of the codebooks (i.e., centroids and labels) so that we could apply or update the codebooks whenever the model weights change.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27


from torch.nn import parameter
class KMeansQuantizer:
 def __init__(self, model : nn.Module, bitwidth=4):
 self.codebook = KMeansQuantizer.quantize(model, bitwidth)

 @torch.no_grad()
 def apply(self, model, update_centroids):
 for name, param in model.named_parameters():
 if name in self.codebook:
 if update_centroids:
 update_codebook(param, codebook=self.codebook[name])
 self.codebook[name] = k_means_quantize(
 param, codebook=self.codebook[name])

 @staticmethod
 @torch.no_grad()
 def quantize(model: nn.Module, bitwidth=4):
 codebook = dict()
 if isinstance(bitwidth, dict):
 for name, param in model.named_parameters():
 if name in bitwidth:
 codebook[name] = k_means_quantize(param, bitwidth=bitwidth[name])
 else:
 for name, param in model.named_parameters():
 if param.dim() > 1:
 codebook[name] = k_means_quantize(param, bitwidth=bitwidth)
 return codebook

Now let’s quantize model into 8 bits, 4 bits and 2 bits using K-Means Quantization. Note that we ignore the storage for codebooks when calculating the model size.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


print('Note that the storage for codebooks is ignored when calculating the model size.')
quantizers = dict()
for bitwidth in [8, 4, 2]:
 recover_model()
 print(f'k-means quantizing model into {bitwidth} bits')
 quantizer = KMeansQuantizer(model, bitwidth)
 quantized_model_size = get_model_size(model, bitwidth)
 print(f" {bitwidth}-bit k-means quantized model has size={quantized_model_size/MiB:.2f} MiB")
 quantized_model_accuracy = evaluate(model, dataloader['test'])
 print(f" {bitwidth}-bit k-means quantized model has accuracy={quantized_model_accuracy:.2f}%")
 quantizers[bitwidth] = quantizer

1
2
3


Note that the storage for codebooks is ignored when calculating the model size.
k-means quantizing model into 8 bits
 8-bit k-means quantized model has size=8.80 MiB

1

eval: 0%| | 0/20 [00:00<?, ?it/s]

1
2
3


 8-bit k-means quantized model has accuracy=92.76%
k-means quantizing model into 4 bits
 4-bit k-means quantized model has size=4.40 MiB

1

eval: 0%| | 0/20 [00:00<?, ?it/s]

1
2
3


 4-bit k-means quantized model has accuracy=79.07%
k-means quantizing model into 2 bits
 2-bit k-means quantized model has size=2.20 MiB

1

eval: 0%| | 0/20 [00:00<?, ?it/s]

1

 2-bit k-means quantized model has accuracy=10.00%

Trained K-Means Quantization

As we can see from the results of last cell, the accuracy significantly drops when quantizing the model into lower bits. Therefore, we have to perform quantization-aware training to recover the accuracy.

During the k-means quantization-aware training, the centroids are also updated, which is proposed in Deep Compression: Compressing Deep Neural Networks With Pruning, Trained Quantization And Huffman Coding.

The gradient for the centroids is calculated as,

$\frac{\partial \mathcal{L} }{\partial C_k} = \sum_{j} \frac{\partial \mathcal{L} }{\partial W_{j}} \frac{\partial W_{j} }{\partial C_k} = \sum_{j} \frac{\partial \mathcal{L} }{\partial W_{j}} \mathbf{1}(I_{j}=k)$

where $\mathcal{L}$ is the loss, $C_k$ is k-th centroid, $I_{j}$ is the label for weight $W_{j}$. $\mathbf{1}()$ is the indicator function, and $\mathbf{1}(I_{j}=k)$ means $1\;\mathrm{if}\;I_{j}=k\;\mathrm{else}\;0$, i.e., $I_{j}==k$.

Here in the lab, for simplicity, we directly update the centroids according to the latest weights:

$C_k = \frac{\sum_{j}W_{j}\mathbf{1}(I_{j}=k)}{\sum_{j}\mathbf{1}(I_{j}=k)}$

Question 3 (10 pts)

Please complete the following codebook update function.

Hint:

The above equation for updating centroids is indeed using the mean of weights in the same cluster to be the updated centroid value.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


def update_codebook(fp32_tensor: torch.Tensor, codebook: Codebook):
 """
 update the centroids in the codebook using updated fp32_tensor
 :param fp32_tensor: [torch.(cuda.)Tensor]
 :param codebook: [Codebook] (the cluster centroids, the cluster label tensor)
 """
 n_clusters = codebook.centroids.numel()
 fp32_tensor = fp32_tensor.view(-1)
 for k in range(n_clusters):
 ############### YOUR CODE STARTS HERE ###############
 # hint: one line of code
 codebook.centroids[k] = fp32_tensor[codebook.labels == k].mean()
 ############### YOUR CODE ENDS HERE #################

Now let’s run the following code cell to finetune the k-means quantized model to recover the accuracy. We will stop finetuning if accuracy drop is less than 0.5.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33


accuracy_drop_threshold = 0.5
quantizers_before_finetune = copy.deepcopy(quantizers)
quantizers_after_finetune = quantizers

for bitwidth in [8, 4, 2]:
 recover_model()
 quantizer = quantizers[bitwidth]
 print(f'k-means quantizing model into {bitwidth} bits')
 quantizer.apply(model, update_centroids=False)
 quantized_model_size = get_model_size(model, bitwidth)
 print(f" {bitwidth}-bit k-means quantized model has size={quantized_model_size/MiB:.2f} MiB")
 quantized_model_accuracy = evaluate(model, dataloader['test'])
 print(f" {bitwidth}-bit k-means quantized model has accuracy={quantized_model_accuracy:.2f}% before quantization-aware training ")
 accuracy_drop = fp32_model_accuracy - quantized_model_accuracy
 if accuracy_drop > accuracy_drop_threshold:
 print(f" Quantization-aware training due to accuracy drop={accuracy_drop:.2f}% is larger than threshold={accuracy_drop_threshold:.2f}%")
 num_finetune_epochs = 5
 optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
 scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, num_finetune_epochs)
 criterion = nn.CrossEntropyLoss()
 best_accuracy = 0
 epoch = num_finetune_epochs
 while accuracy_drop > accuracy_drop_threshold and epoch > 0:
 train(model, dataloader['train'], criterion, optimizer, scheduler,
 callbacks=[lambda: quantizer.apply(model, update_centroids=True)])
 model_accuracy = evaluate(model, dataloader['test'])
 is_best = model_accuracy > best_accuracy
 best_accuracy = max(model_accuracy, best_accuracy)
 print(f' Epoch {num_finetune_epochs-epoch} Accuracy {model_accuracy:.2f}% / Best Accuracy: {best_accuracy:.2f}%')
 accuracy_drop = fp32_model_accuracy - best_accuracy
 epoch -= 1
 else:
 print(f" No need for quantization-aware training since accuracy drop={accuracy_drop:.2f}% is smaller than threshold={accuracy_drop_threshold:.2f}%")

1
2


k-means quantizing model into 8 bits
 8-bit k-means quantized model has size=8.80 MiB

1

eval: 0%| | 0/20 [00:00<?, ?it/s]

1
2
3
4


 8-bit k-means quantized model has accuracy=92.76% before quantization-aware training 
 No need for quantization-aware training since accuracy drop=0.19% is smaller than threshold=0.50%
k-means quantizing model into 4 bits
 4-bit k-means quantized model has size=4.40 MiB

1

eval: 0%| | 0/20 [00:00<?, ?it/s]

1
2


 4-bit k-means quantized model has accuracy=79.07% before quantization-aware training 
 Quantization-aware training due to accuracy drop=13.88% is larger than threshold=0.50%

1

train: 0%| | 0/98 [00:00<?, ?it/s]

1

eval: 0%| | 0/20 [00:00<?, ?it/s]

1
2
3


 Epoch 0 Accuracy 92.47% / Best Accuracy: 92.47%
k-means quantizing model into 2 bits
 2-bit k-means quantized model has size=2.20 MiB

1

eval: 0%| | 0/20 [00:00<?, ?it/s]

1
2


 2-bit k-means quantized model has accuracy=10.00% before quantization-aware training 
 Quantization-aware training due to accuracy drop=82.95% is larger than threshold=0.50%

1

train: 0%| | 0/98 [00:00<?, ?it/s]

1

eval: 0%| | 0/20 [00:00<?, ?it/s]

1

 Epoch 0 Accuracy 90.21% / Best Accuracy: 90.21%

1

train: 0%| | 0/98 [00:00<?, ?it/s]

1

eval: 0%| | 0/20 [00:00<?, ?it/s]

1

 Epoch 1 Accuracy 90.82% / Best Accuracy: 90.82%

1

train: 0%| | 0/98 [00:00<?, ?it/s]

1

eval: 0%| | 0/20 [00:00<?, ?it/s]

1

 Epoch 2 Accuracy 91.00% / Best Accuracy: 91.00%

1

train: 0%| | 0/98 [00:00<?, ?it/s]

1

eval: 0%| | 0/20 [00:00<?, ?it/s]

1

 Epoch 3 Accuracy 91.13% / Best Accuracy: 91.13%

1

train: 0%| | 0/98 [00:00<?, ?it/s]

1

eval: 0%| | 0/20 [00:00<?, ?it/s]

1

 Epoch 4 Accuracy 91.20% / Best Accuracy: 91.20%

Linear Quantization

In this section, we will implement and perform linear quantization.

Linear quantization directly rounds the floating-point value into the nearest quantized integer after range truncation and scaling.

Linear quantization can be represented as

$r = S(q-Z)$

where $r$ is a floating point real number, $q$ is a n-bit integer, $Z$ is a n-bit integer, and $S$ is a floating point real number. $Z$ is quantization zero point and $S$ is quantization scaling factor. Both constant $Z$ and $S$ are quantization parameters.

n-bit Integer

A n-bit signed integer is usually represented in two’s complement notation.

A n-bit signed integer can enode integers in the range $[-2^{n-1}, 2^{n-1}-1]$. For example, a 8-bit integer falls in the range [-128, 127].

1
2
3
4


def get_quantized_range(bitwidth):
 quantized_max = (1 << (bitwidth - 1)) - 1
 quantized_min = -(1 << (bitwidth - 1))
 return quantized_min, quantized_max

Question 4 (10 pts)

Please complete the following linear quantization function.

Hint:

From $r=S(q-Z)$, we have $q = r/S + Z$.
Both $r$ and $S$ are floating numbers, and thus we cannot directly add integer $Z$ to $r/S$. Therefore $q = \mathrm{int}(\mathrm{round}(r/S)) + Z$.
To convert torch.FloatTensor to torch.IntTensor, we could use torch.round(), torch.Tensor.round(), torch.Tensor.round_() to first convert all values to floating integer, and then use torch.Tensor.to(torch.int8) to convert the data type from torch.float to torch.int8.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38


def linear_quantize(fp_tensor, bitwidth, scale, zero_point, dtype=torch.int8) -> torch.Tensor:
 """
 linear quantization for single fp_tensor
 from
 fp_tensor = (quantized_tensor - zero_point) * scale
 we have,
 quantized_tensor = int(round(fp_tensor / scale)) + zero_point
 :param tensor: [torch.(cuda.)FloatTensor] floating tensor to be quantized
 :param bitwidth: [int] quantization bit width
 :param scale: [torch.(cuda.)FloatTensor] scaling factor
 :param zero_point: [torch.(cuda.)IntTensor] the desired centroid of tensor values
 :return:
 [torch.(cuda.)FloatTensor] quantized tensor whose values are integers
 """
 assert(fp_tensor.dtype == torch.float)
 assert(isinstance(scale, float) or
 (scale.dtype == torch.float and scale.dim() == fp_tensor.dim()))
 assert(isinstance(zero_point, int) or
 (zero_point.dtype == dtype and zero_point.dim() == fp_tensor.dim()))

 ############### YOUR CODE STARTS HERE ###############
 # Step 1: scale the fp_tensor
 scaled_tensor = fp_tensor / scale
 # Step 2: round the floating value to integer value
 rounded_tensor = scaled_tensor.round()
 ############### YOUR CODE ENDS HERE #################

 rounded_tensor = rounded_tensor.to(dtype)

 ############### YOUR CODE STARTS HERE ###############
 # Step 3: shift the rounded_tensor to make zero_point 0
 shifted_tensor = rounded_tensor + zero_point
 ############### YOUR CODE ENDS HERE #################

 # Step 4: clamp the shifted_tensor to lie in bitwidth-bit range
 quantized_min, quantized_max = get_quantized_range(bitwidth)
 quantized_tensor = shifted_tensor.clamp_(quantized_min, quantized_max)
 return quantized_tensor

Let’s verify the functionality of defined linear quantization by applying the function above on a dummy tensor.

1

test_linear_quantize()

1
2
3
4
5


* Test linear_quantize()
 target bitwidth: 2 bits
 scale: 0.3333333333333333
 zero point: -1
* Test passed.

1

<Figure size 1000x3200 with 3 Axes>

Question 5 (15 pts)

Now we have to determine the scaling factor $S$ and zero point $Z$ for linear quantization.

Recall that linear quantization can be represented as

$r = S(q-Z)$

Scale

Linear quantization projects the floating point range [fp_min, fp_max] to the quantized range [quantized_min, quantized_max]. That is to say,

$r_{\mathrm{max}} = S(q_{\mathrm{max}}-Z)$

$r_{\mathrm{min}} = S(q_{\mathrm{min}}-Z)$

Substracting these two equations, we have,

Question 5.1 (3 pts)

Please select the correct answer and delete the wrong answers in the next text cell.

$S=(r_{\mathrm{max}} - r_{\mathrm{min}}) / (q_{\mathrm{max}} - q_{\mathrm{min}})$

There are different approaches to determine the $r_{\mathrm{min}}$ and $r_{\mathrm{max}}$ of a floating point tensor fp_tensor.

The most common method is directly using the minimum and maximum value of fp_tensor.
Another widely used method is minimizing Kullback-Leibler-J divergence to determine the fp_max.

zero point

Once we determine the scaling factor $S$, we can directly use the relationship between $r_{\mathrm{min}}$ and $q_{\mathrm{min}}$ to calculate the zero point $Z$.

Question 5.2 (4 pts)

Please select the correct answer and delete the wrong answers in the next text cell.

$Z = \mathrm{int}(\mathrm{round}(q_{\mathrm{min}} - r_{\mathrm{min}} / S))$

Question 5.3 (8 pts)

Please complete the following function for calculating the scale $S$ and zero point $Z$ from floating point tensor $r$.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28


def get_quantization_scale_and_zero_point(fp_tensor, bitwidth):
 """
 get quantization scale for single tensor
 :param fp_tensor: [torch.(cuda.)Tensor] floating tensor to be quantized
 :param bitwidth: [int] quantization bit width
 :return:
 [float] scale
 [int] zero_point
 """
 quantized_min, quantized_max = get_quantized_range(bitwidth)
 fp_max = fp_tensor.max().item()
 fp_min = fp_tensor.min().item()

 ############### YOUR CODE STARTS HERE ###############
 # hint: one line of code for calculating scale
 scale = (fp_max - fp_min) / (quantized_max - quantized_min)
 # hint: one line of code for calculating zero_point
 zero_point = quantized_min - fp_min / scale
 ############### YOUR CODE ENDS HERE #################

 # clip the zero_point to fall in [quantized_min, quantized_max]
 if zero_point < quantized_min:
 zero_point = quantized_min
 elif zero_point > quantized_max:
 zero_point = quantized_max
 else: # convert from float to int using round()
 zero_point = round(zero_point)
 return scale, int(zero_point)

We now wrap linear_quantize() in Question 4 and get_quantization_scale_and_zero_point() in Question 5 into one function.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


def linear_quantize_feature(fp_tensor, bitwidth):
 """
 linear quantization for feature tensor
 :param fp_tensor: [torch.(cuda.)Tensor] floating feature to be quantized
 :param bitwidth: [int] quantization bit width
 :return:
 [torch.(cuda.)Tensor] quantized tensor
 [float] scale tensor
 [int] zero point
 """
 scale, zero_point = get_quantization_scale_and_zero_point(fp_tensor, bitwidth)
 quantized_tensor = linear_quantize(fp_tensor, bitwidth, scale, zero_point)
 return quantized_tensor, scale, zero_point

Special case: linear quantization on weight tensor

Let’s first see the distribution of weight values.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31


def plot_weight_distribution(model, bitwidth=32):
 # bins = (1 << bitwidth) if bitwidth <= 8 else 256
 if bitwidth <= 8:
 qmin, qmax = get_quantized_range(bitwidth)
 bins = np.arange(qmin, qmax + 2)
 align = 'left'
 else:
 bins = 256
 align = 'mid'
 fig, axes = plt.subplots(3,3, figsize=(10, 6))
 axes = axes.ravel()
 plot_index = 0
 for name, param in model.named_parameters():
 if param.dim() > 1:
 ax = axes[plot_index]
 ax.hist(param.detach().view(-1).cpu(), bins=bins, density=True,
 align=align, color = 'blue', alpha = 0.5,
 edgecolor='black' if bitwidth <= 4 else None)
 if bitwidth <= 4:
 quantized_min, quantized_max = get_quantized_range(bitwidth)
 ax.set_xticks(np.arange(start=quantized_min, stop=quantized_max+1))
 ax.set_xlabel(name)
 ax.set_ylabel('density')
 plot_index += 1
 fig.suptitle(f'Histogram of Weights (bitwidth={bitwidth} bits)')
 fig.tight_layout()
 fig.subplots_adjust(top=0.925)
 plt.show()

recover_model()
plot_weight_distribution(model)

1

<Figure size 1000x600 with 9 Axes>

As we can see from the histograms above, the distribution of weight values are nearly symmetric about 0 (except for the classifier in this case). Therefore, we usually make zero point $Z=0$ when quantizating the weights.

From $r = S(q-Z)$, we have

$r_{\mathrm{max}} = S \cdot q_{\mathrm{max}}$

and then

$S = r_{\mathrm{max}} / q_{\mathrm{max}}$

We directly use the maximum magnitude of weight values as $r_{\mathrm{max}}$.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


def get_quantization_scale_for_weight(weight, bitwidth):
 """
 get quantization scale for single tensor of weight
 :param weight: [torch.(cuda.)Tensor] floating weight to be quantized
 :param bitwidth: [integer] quantization bit width
 :return:
 [floating scalar] scale
 """
 # we just assume values in weight are symmetric
 # we also always make zero_point 0 for weight
 fp_max = max(weight.abs().max().item(), 5e-7)
 _, quantized_max = get_quantized_range(bitwidth)
 return fp_max / quantized_max

Per-channel Linear Quantization

Recall that for 2D convolution, the weight tensor is a 4-D tensor in the shape of (num_output_channels, num_input_channels, kernel_height, kernel_width).

Intensive experiments show that using the different scaling factors $S$ and zero points $Z$ for different output channels will perform better. Therefore, we have to determine scaling factor $S$ and zero point $Z$ for the subtensor of each output channel independently.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23


def linear_quantize_weight_per_channel(tensor, bitwidth):
 """
 linear quantization for weight tensor
 using different scales and zero_points for different output channels
 :param tensor: [torch.(cuda.)Tensor] floating weight to be quantized
 :param bitwidth: [int] quantization bit width
 :return:
 [torch.(cuda.)Tensor] quantized tensor
 [torch.(cuda.)Tensor] scale tensor
 [int] zero point (which is always 0)
 """
 dim_output_channels = 0
 num_output_channels = tensor.shape[dim_output_channels]
 scale = torch.zeros(num_output_channels, device=tensor.device)
 for oc in range(num_output_channels):
 _subtensor = tensor.select(dim_output_channels, oc)
 _scale = get_quantization_scale_for_weight(_subtensor, bitwidth)
 scale[oc] = _scale
 scale_shape = [1] * tensor.dim()
 scale_shape[dim_output_channels] = -1
 scale = scale.view(scale_shape)
 quantized_tensor = linear_quantize(tensor, bitwidth, scale, zero_point=0)
 return quantized_tensor, scale, 0

A Quick Peek at Linear Quantization on Weights

Now let’s have a peek on the weight distribution and model size when applying linear quantization on weights with different bitwidths.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12


@torch.no_grad()
def peek_linear_quantization():
 for bitwidth in [4, 2]:
 for name, param in model.named_parameters():
 if param.dim() > 1:
 quantized_param, scale, zero_point = \
 linear_quantize_weight_per_channel(param, bitwidth)
 param.copy_(quantized_param)
 plot_weight_distribution(model, bitwidth)
 recover_model()

peek_linear_quantization()

1

<Figure size 1000x600 with 9 Axes>

1

<Figure size 1000x600 with 9 Axes>

Quantized Inference

After quantization, the inference of convolution and fully-connected layers also change.

Recall that $r = S(q-Z)$, and we have

$r_{\mathrm{input}} = S_{\mathrm{input}}(q_{\mathrm{input}}-Z_{\mathrm{input}})$

$r_{\mathrm{weight}} = S_{\mathrm{weight}}(q_{\mathrm{weight}}-Z_{\mathrm{weight}})$

$r_{\mathrm{bias}} = S_{\mathrm{bias}}(q_{\mathrm{bias}}-Z_{\mathrm{bias}})$

Since $Z_{\mathrm{weight}}=0$, $r_{\mathrm{weight}} = S_{\mathrm{weight}}q_{\mathrm{weight}}$.

The floating point convolution can be written as,

$r_{\mathrm{output}} = \mathrm{CONV}[r_{\mathrm{input}}, r_{\mathrm{weight}}] + r_{\mathrm{bias}}\\ \;\;\;\;\;\;\;\;= \mathrm{CONV}[S_{\mathrm{input}}(q_{\mathrm{input}}-Z_{\mathrm{input}}), S_{\mathrm{weight}}q_{\mathrm{weight}}] + S_{\mathrm{bias}}(q_{\mathrm{bias}}-Z_{\mathrm{bias}})\\ \;\;\;\;\;\;\;\;= \mathrm{CONV}[q_{\mathrm{input}}-Z_{\mathrm{input}}, q_{\mathrm{weight}}]\cdot (S_{\mathrm{input}} \cdot S_{\mathrm{weight}}) + S_{\mathrm{bias}}(q_{\mathrm{bias}}-Z_{\mathrm{bias}})$

To further simplify the computation, we could let

$Z_{\mathrm{bias}} = 0$

$S_{\mathrm{bias}} = S_{\mathrm{input}} \cdot S_{\mathrm{weight}}$

so that

$r_{\mathrm{output}} = (\mathrm{CONV}[q_{\mathrm{input}}-Z_{\mathrm{input}}, q_{\mathrm{weight}}] + q_{\mathrm{bias}})\cdot (S_{\mathrm{input}} \cdot S_{\mathrm{weight}})$ $\;\;\;\;\;\;\;\;= (\mathrm{CONV}[q_{\mathrm{input}}, q_{\mathrm{weight}}] - \mathrm{CONV}[Z_{\mathrm{input}}, q_{\mathrm{weight}}] + q_{\mathrm{bias}})\cdot (S_{\mathrm{input}}S_{\mathrm{weight}})$

Since

$r_{\mathrm{output}} = S_{\mathrm{output}}(q_{\mathrm{output}}-Z_{\mathrm{output}})$

we have

$S_{\mathrm{output}}(q_{\mathrm{output}}-Z_{\mathrm{output}}) = (\mathrm{CONV}[q_{\mathrm{input}}, q_{\mathrm{weight}}] - \mathrm{CONV}[Z_{\mathrm{input}}, q_{\mathrm{weight}}] + q_{\mathrm{bias}})\cdot (S_{\mathrm{input}} S_{\mathrm{weight}})$

and thus

$q_{\mathrm{output}} = (\mathrm{CONV}[q_{\mathrm{input}}, q_{\mathrm{weight}}] - \mathrm{CONV}[Z_{\mathrm{input}}, q_{\mathrm{weight}}] + q_{\mathrm{bias}})\cdot (S_{\mathrm{input}}S_{\mathrm{weight}} / S_{\mathrm{output}}) + Z_{\mathrm{output}}$

Since $Z_{\mathrm{input}}$, $q_{\mathrm{weight}}$, $q_{\mathrm{bias}}$ are determined before inference, let

$Q_{\mathrm{bias}} = q_{\mathrm{bias}} - \mathrm{CONV}[Z_{\mathrm{input}}, q_{\mathrm{weight}}]$

we have

$q_{\mathrm{output}} = (\mathrm{CONV}[q_{\mathrm{input}}, q_{\mathrm{weight}}] + Q_{\mathrm{bias}}) \cdot (S_{\mathrm{input}}S_{\mathrm{weight}} / S_{\mathrm{output}}) + Z_{\mathrm{output}}$

Similarily, for fully-connected layer, we have

$q_{\mathrm{output}} = (\mathrm{Linear}[q_{\mathrm{input}}, q_{\mathrm{weight}}] + Q_{\mathrm{bias}})\cdot (S_{\mathrm{input}} \cdot S_{\mathrm{weight}} / S_{\mathrm{output}}) + Z_{\mathrm{output}}$

where

$Q_{\mathrm{bias}} = q_{\mathrm{bias}} - \mathrm{Linear}[Z_{\mathrm{input}}, q_{\mathrm{weight}}]$

Question 6 (5 pts)

Please complete the following function for linear quantizing the bias.

Hint:

From the above deduction, we know that

$Z_{\mathrm{bias}} = 0$

$S_{\mathrm{bias}} = S_{\mathrm{input}} \cdot S_{\mathrm{weight}}$

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26


def linear_quantize_bias_per_output_channel(bias, weight_scale, input_scale):
 """
 linear quantization for single bias tensor
 quantized_bias = fp_bias / bias_scale
 :param bias: [torch.FloatTensor] bias weight to be quantized
 :param weight_scale: [float or torch.FloatTensor] weight scale tensor
 :param input_scale: [float] input scale
 :return:
 [torch.IntTensor] quantized bias tensor
 """
 assert(bias.dim() == 1)
 assert(bias.dtype == torch.float)
 assert(isinstance(input_scale, float))
 if isinstance(weight_scale, torch.Tensor):
 assert(weight_scale.dtype == torch.float)
 weight_scale = weight_scale.view(-1)
 assert(bias.numel() == weight_scale.numel())

 ############### YOUR CODE STARTS HERE ###############
 # hint: one line of code
 bias_scale = input_scale * weight_scale
 ############### YOUR CODE ENDS HERE #################

 quantized_bias = linear_quantize(bias, 32, bias_scale,
 zero_point=0, dtype=torch.int32)
 return quantized_bias, bias_scale, 0

Quantized Fully-Connected Layer

For quantized fully-connected layer, we first precompute $Q_{\mathrm{bias}}$. Recall that $Q_{\mathrm{bias}} = q_{\mathrm{bias}} - \mathrm{Linear}[Z_{\mathrm{input}}, q_{\mathrm{weight}}]$.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


def shift_quantized_linear_bias(quantized_bias, quantized_weight, input_zero_point):
 """
 shift quantized bias to incorporate input_zero_point for nn.Linear
 shifted_quantized_bias = quantized_bias - Linear(input_zero_point, quantized_weight)
 :param quantized_bias: [torch.IntTensor] quantized bias (torch.int32)
 :param quantized_weight: [torch.CharTensor] quantized weight (torch.int8)
 :param input_zero_point: [int] input zero point
 :return:
 [torch.IntTensor] shifted quantized bias tensor
 """
 assert(quantized_bias.dtype == torch.int32)
 assert(isinstance(input_zero_point, int))
 return quantized_bias - quantized_weight.sum(1).to(torch.int32) * input_zero_point

Question 7 (15 pts)

Please complete the following quantized fully-connected layer inference function.

Hint:

$q_{\mathrm{output}} = (\mathrm{Linear}[q_{\mathrm{input}}, q_{\mathrm{weight}}] + Q_{\mathrm{bias}})\cdot (S_{\mathrm{input}} S_{\mathrm{weight}} / S_{\mathrm{output}}) + Z_{\mathrm{output}}$

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49


def quantized_linear(input, weight, bias, feature_bitwidth, weight_bitwidth,
 input_zero_point, output_zero_point,
 input_scale, weight_scale, output_scale):
 """
 quantized fully-connected layer
 :param input: [torch.CharTensor] quantized input (torch.int8)
 :param weight: [torch.CharTensor] quantized weight (torch.int8)
 :param bias: [torch.IntTensor] shifted quantized bias or None (torch.int32)
 :param feature_bitwidth: [int] quantization bit width of input and output
 :param weight_bitwidth: [int] quantization bit width of weight
 :param input_zero_point: [int] input zero point
 :param output_zero_point: [int] output zero point
 :param input_scale: [float] input feature scale
 :param weight_scale: [torch.FloatTensor] weight per-channel scale
 :param output_scale: [float] output feature scale
 :return:
 [torch.CharIntTensor] quantized output feature (torch.int8)
 """
 assert(input.dtype == torch.int8)
 assert(weight.dtype == input.dtype)
 assert(bias is None or bias.dtype == torch.int32)
 assert(isinstance(input_zero_point, int))
 assert(isinstance(output_zero_point, int))
 assert(isinstance(input_scale, float))
 assert(isinstance(output_scale, float))
 assert(weight_scale.dtype == torch.float)

 # Step 1: integer-based fully-connected (8-bit multiplication with 32-bit accumulation)
 if 'cpu' in input.device.type:
 # use 32-b MAC for simplicity
 output = torch.nn.functional.linear(input.to(torch.int32), weight.to(torch.int32), bias)
 else:
 # current version pytorch does not yet support integer-based linear() on GPUs
 output = torch.nn.functional.linear(input.float(), weight.float(), bias.float())

 ############### YOUR CODE STARTS HERE ###############
 # Step 2: scale the output
 # hint: 1. scales are floating numbers, we need to convert output to float as well
 # 2. the shape of weight scale is [oc, 1, 1, 1] while the shape of output is [batch_size, oc]
 output = output.float() * (input_scale * weight_scale / output_scale).view(1, -1)

 # Step 3: shift output by output_zero_point
 # hint: one line of code
 output = output + output_zero_point
 ############### YOUR CODE ENDS HERE #################

 # Make sure all value lies in the bitwidth-bit range
 output = output.round().clamp(*get_quantized_range(feature_bitwidth)).to(torch.int8)
 return output

Let’s verify the functionality of defined quantized fully connected layer.

1

test_quantized_fc()

1
2
3
4
5
6


* Test quantized_fc()
 target bitwidth: 2 bits
 batch size: 4
 input channels: 8
 output channels: 8
* Test passed.

1

<Figure size 1500x1200 with 9 Axes>

Quantized Convolution

For quantized convolution layer, we first precompute $Q_{\mathrm{bias}}$. Recall that $Q_{\mathrm{bias}} = q_{\mathrm{bias}} - \mathrm{CONV}[Z_{\mathrm{input}}, q_{\mathrm{weight}}]$.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


def shift_quantized_conv2d_bias(quantized_bias, quantized_weight, input_zero_point):
 """
 shift quantized bias to incorporate input_zero_point for nn.Conv2d
 shifted_quantized_bias = quantized_bias - Conv(input_zero_point, quantized_weight)
 :param quantized_bias: [torch.IntTensor] quantized bias (torch.int32)
 :param quantized_weight: [torch.CharTensor] quantized weight (torch.int8)
 :param input_zero_point: [int] input zero point
 :return:
 [torch.IntTensor] shifted quantized bias tensor
 """
 assert(quantized_bias.dtype == torch.int32)
 assert(isinstance(input_zero_point, int))
 return quantized_bias - quantized_weight.sum((1,2,3)).to(torch.int32) * input_zero_point

Question 8 (10 pts)

Please complete the following quantized convolution function.

Hint:

$q_{\mathrm{output}} = (\mathrm{CONV}[q_{\mathrm{input}}, q_{\mathrm{weight}}] + Q_{\mathrm{bias}}) \cdot (S_{\mathrm{input}}S_{\mathrm{weight}} / S_{\mathrm{output}}) + Z_{\mathrm{output}}$

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57


def quantized_conv2d(input, weight, bias, feature_bitwidth, weight_bitwidth,
 input_zero_point, output_zero_point,
 input_scale, weight_scale, output_scale,
 stride, padding, dilation, groups):
 """
 quantized 2d convolution
 :param input: [torch.CharTensor] quantized input (torch.int8)
 :param weight: [torch.CharTensor] quantized weight (torch.int8)
 :param bias: [torch.IntTensor] shifted quantized bias or None (torch.int32)
 :param feature_bitwidth: [int] quantization bit width of input and output
 :param weight_bitwidth: [int] quantization bit width of weight
 :param input_zero_point: [int] input zero point
 :param output_zero_point: [int] output zero point
 :param input_scale: [float] input feature scale
 :param weight_scale: [torch.FloatTensor] weight per-channel scale
 :param output_scale: [float] output feature scale
 :return:
 [torch.(cuda.)CharTensor] quantized output feature
 """
 assert(len(padding) == 4)
 assert(input.dtype == torch.int8)
 assert(weight.dtype == input.dtype)
 assert(bias is None or bias.dtype == torch.int32)
 assert(isinstance(input_zero_point, int))
 assert(isinstance(output_zero_point, int))
 assert(isinstance(input_scale, float))
 assert(isinstance(output_scale, float))
 assert(weight_scale.dtype == torch.float)

 # Step 1: calculate integer-based 2d convolution (8-bit multiplication with 32-bit accumulation)
 input = torch.nn.functional.pad(input, padding, 'constant', input_zero_point)
 if 'cpu' in input.device.type:
 # use 32-b MAC for simplicity
 output = torch.nn.functional.conv2d(input.to(torch.int32), weight.to(torch.int32), None, stride, 0, dilation, groups)
 else:
 # current version pytorch does not yet support integer-based conv2d() on GPUs
 output = torch.nn.functional.conv2d(input.float(), weight.float(), None, stride, 0, dilation, groups)
 output = output.round().to(torch.int32)
 if bias is not None:
 output = output + bias.view(1, -1, 1, 1)

 ############### YOUR CODE STARTS HERE ###############
 # hint: this code block should be the very similar to quantized_linear()

 # Step 2: scale the output
 # hint: 1. scales are floating numbers, we need to convert output to float as well
 # 2. the shape of weight scale is [oc, 1, 1, 1] while the shape of output is [batch_size, oc, height, width]
 output = output.float() * (input_scale * weight_scale / output_scale).view(1, -1, 1, 1)

 # Step 3: shift output by output_zero_point
 # hint: one line of code
 output = output + output_zero_point
 ############### YOUR CODE ENDS HERE #################

 # Make sure all value lies in the bitwidth-bit range
 output = output.round().clamp(*get_quantized_range(feature_bitwidth)).to(torch.int8)
 return output

Question 9 (10 pts)

Finally, we are putting everything together and perform post-training int8 quantization for the model. We will convert the convolutional and linear layers in the model to a quantized version one-by-one.

Firstly, we will fuse a BatchNorm layer into its previous convolutional layer, which is a standard practice before quantization. Fusing batchnorm reduces the extra multiplication during inference.

We will also verify that the fused model model_fused has the same accuracy as the original model (BN fusion is an equivalent transform that does not change network functionality).

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35


def fuse_conv_bn(conv, bn):
 # modified from https://mmcv.readthedocs.io/en/latest/_modules/mmcv/cnn/utils/fuse_conv_bn.html
 assert conv.bias is None

 factor = bn.weight.data / torch.sqrt(bn.running_var.data + bn.eps)
 conv.weight.data = conv.weight.data * factor.reshape(-1, 1, 1, 1)
 conv.bias = nn.Parameter(- bn.running_mean.data * factor + bn.bias.data)

 return conv

print('Before conv-bn fusion: backbone length', len(model.backbone))
# fuse the batchnorm into conv layers
recover_model()
model_fused = copy.deepcopy(model)
fused_backbone = []
ptr = 0
while ptr < len(model_fused.backbone):
 if isinstance(model_fused.backbone[ptr], nn.Conv2d) and \
 isinstance(model_fused.backbone[ptr + 1], nn.BatchNorm2d):
 fused_backbone.append(fuse_conv_bn(
 model_fused.backbone[ptr], model_fused.backbone[ptr+ 1]))
 ptr += 2
 else:
 fused_backbone.append(model_fused.backbone[ptr])
 ptr += 1
model_fused.backbone = nn.Sequential(*fused_backbone)

print('After conv-bn fusion: backbone length', len(model_fused.backbone))
# sanity check, no BN anymore
for m in model_fused.modules():
 assert not isinstance(m, nn.BatchNorm2d)

# the accuracy will remain the same after fusion
fused_acc = evaluate(model_fused, dataloader['test'])
print(f'Accuracy of the fused model={fused_acc:.2f}%')

1
2


Before conv-bn fusion: backbone length 29
After conv-bn fusion: backbone length 21

1

eval: 0%| | 0/20 [00:00<?, ?it/s]

1

Accuracy of the fused model=92.95%

We will run the model with some sample data to get the range of each feature map, so that we can get the range of the feature maps and compute their corresponding scaling factors and zero points.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25


# add hook to record the min max value of the activation
input_activation = {}
output_activation = {}

def add_range_recoder_hook(model):
 import functools
 def _record_range(self, x, y, module_name):
 x = x[0]
 input_activation[module_name] = x.detach()
 output_activation[module_name] = y.detach()

 all_hooks = []
 for name, m in model.named_modules():
 if isinstance(m, (nn.Conv2d, nn.Linear, nn.ReLU)):
 all_hooks.append(m.register_forward_hook(
 functools.partial(_record_range, module_name=name)))
 return all_hooks

hooks = add_range_recoder_hook(model_fused)
sample_data = iter(dataloader['train']).__next__()[0]
model_fused(sample_data.cuda())

# remove hooks
for h in hooks:
 h.remove()

Finally, let’s do model quantization. We will convert the model in the following mapping

1
2
3
4
5
6
7


nn.Conv2d: QuantizedConv2d,
nn.Linear: QuantizedLinear,
# the following twos are just wrappers, as current
# torch modules do not support int8 data format;
# we will temporarily convert them to fp32 for computation
nn.MaxPool2d: QuantizedMaxPool2d,
nn.AvgPool2d: QuantizedAvgPool2d,

 1
 2
 3
 4
 5
 6
 7
 8
 9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156


class QuantizedConv2d(nn.Module):
 def __init__(self, weight, bias,
 input_zero_point, output_zero_point,
 input_scale, weight_scale, output_scale,
 stride, padding, dilation, groups,
 feature_bitwidth=8, weight_bitwidth=8):
 super().__init__()
 # current version Pytorch does not support IntTensor as nn.Parameter
 self.register_buffer('weight', weight)
 self.register_buffer('bias', bias)

 self.input_zero_point = input_zero_point
 self.output_zero_point = output_zero_point

 self.input_scale = input_scale
 self.register_buffer('weight_scale', weight_scale)
 self.output_scale = output_scale

 self.stride = stride
 self.padding = (padding[1], padding[1], padding[0], padding[0])
 self.dilation = dilation
 self.groups = groups

 self.feature_bitwidth = feature_bitwidth
 self.weight_bitwidth = weight_bitwidth


 def forward(self, x):
 return quantized_conv2d(
 x, self.weight, self.bias,
 self.feature_bitwidth, self.weight_bitwidth,
 self.input_zero_point, self.output_zero_point,
 self.input_scale, self.weight_scale, self.output_scale,
 self.stride, self.padding, self.dilation, self.groups
 )

class QuantizedLinear(nn.Module):
 def __init__(self, weight, bias,
 input_zero_point, output_zero_point,
 input_scale, weight_scale, output_scale,
 feature_bitwidth=8, weight_bitwidth=8):
 super().__init__()
 # current version Pytorch does not support IntTensor as nn.Parameter
 self.register_buffer('weight', weight)
 self.register_buffer('bias', bias)

 self.input_zero_point = input_zero_point
 self.output_zero_point = output_zero_point

 self.input_scale = input_scale
 self.register_buffer('weight_scale', weight_scale)
 self.output_scale = output_scale

 self.feature_bitwidth = feature_bitwidth
 self.weight_bitwidth = weight_bitwidth

 def forward(self, x):
 return quantized_linear(
 x, self.weight, self.bias,
 self.feature_bitwidth, self.weight_bitwidth,
 self.input_zero_point, self.output_zero_point,
 self.input_scale, self.weight_scale, self.output_scale
 )

class QuantizedMaxPool2d(nn.MaxPool2d):
 def forward(self, x):
 # current version PyTorch does not support integer-based MaxPool
 return super().forward(x.float()).to(torch.int8)

class QuantizedAvgPool2d(nn.AvgPool2d):
 def forward(self, x):
 # current version PyTorch does not support integer-based AvgPool
 return super().forward(x.float()).to(torch.int8)

# we use int8 quantization, which is quite popular
feature_bitwidth = weight_bitwidth = 8
quantized_model = copy.deepcopy(model_fused)
quantized_backbone = []
ptr = 0
while ptr < len(quantized_model.backbone):
 if isinstance(quantized_model.backbone[ptr], nn.Conv2d) and \
 isinstance(quantized_model.backbone[ptr + 1], nn.ReLU):
 conv = quantized_model.backbone[ptr]
 conv_name = f'backbone.{ptr}'
 relu = quantized_model.backbone[ptr + 1]
 relu_name = f'backbone.{ptr + 1}'

 input_scale, input_zero_point = \
 get_quantization_scale_and_zero_point(
 input_activation[conv_name], feature_bitwidth)

 output_scale, output_zero_point = \
 get_quantization_scale_and_zero_point(
 output_activation[relu_name], feature_bitwidth)

 quantized_weight, weight_scale, weight_zero_point = \
 linear_quantize_weight_per_channel(conv.weight.data, weight_bitwidth)
 quantized_bias, bias_scale, bias_zero_point = \
 linear_quantize_bias_per_output_channel(
 conv.bias.data, weight_scale, input_scale)
 shifted_quantized_bias = \
 shift_quantized_conv2d_bias(quantized_bias, quantized_weight,
 input_zero_point)

 quantized_conv = QuantizedConv2d(
 quantized_weight, shifted_quantized_bias,
 input_zero_point, output_zero_point,
 input_scale, weight_scale, output_scale,
 conv.stride, conv.padding, conv.dilation, conv.groups,
 feature_bitwidth=feature_bitwidth, weight_bitwidth=weight_bitwidth
 )

 quantized_backbone.append(quantized_conv)
 ptr += 2
 elif isinstance(quantized_model.backbone[ptr], nn.MaxPool2d):
 quantized_backbone.append(QuantizedMaxPool2d(
 kernel_size=quantized_model.backbone[ptr].kernel_size,
 stride=quantized_model.backbone[ptr].stride
 ))
 ptr += 1
 elif isinstance(quantized_model.backbone[ptr], nn.AvgPool2d):
 quantized_backbone.append(QuantizedAvgPool2d(
 kernel_size=quantized_model.backbone[ptr].kernel_size,
 stride=quantized_model.backbone[ptr].stride
 ))
 ptr += 1
 else:
 raise NotImplementedError(type(quantized_model.backbone[ptr])) # should not happen
quantized_model.backbone = nn.Sequential(*quantized_backbone)

# finally, quantized the classifier
fc_name = 'classifier'
fc = model.classifier
input_scale, input_zero_point = \
 get_quantization_scale_and_zero_point(
 input_activation[fc_name], feature_bitwidth)

output_scale, output_zero_point = \
 get_quantization_scale_and_zero_point(
 output_activation[fc_name], feature_bitwidth)

quantized_weight, weight_scale, weight_zero_point = \
 linear_quantize_weight_per_channel(fc.weight.data, weight_bitwidth)
quantized_bias, bias_scale, bias_zero_point = \
 linear_quantize_bias_per_output_channel(
 fc.bias.data, weight_scale, input_scale)
shifted_quantized_bias = \
 shift_quantized_linear_bias(quantized_bias, quantized_weight,
 input_zero_point)

quantized_model.classifier = QuantizedLinear(
 quantized_weight, shifted_quantized_bias,
 input_zero_point, output_zero_point,
 input_scale, weight_scale, output_scale,
 feature_bitwidth=feature_bitwidth, weight_bitwidth=weight_bitwidth
)

The quantization process is done! Let’s print and visualize the model architecture and also verify the accuracy of the quantized model.

Question 9.1 (5 pts)

To run the quantized model, we need an extra preprocessing to map the input data from range (0, 1) into int8 range of (-128, 127). Fill in the code below to finish the extra preprocessing.

Hint: you should find that the quantized model has roughly the same accuracy as the fp32 counterpart.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12


print(quantized_model)

def extra_preprocess(x):
 # hint: you need to convert the original fp32 input of range (0, 1)
 # into int8 format of range (-128, 127)
 ############### YOUR CODE STARTS HERE ###############
 return (x * 255 - 128).round().clamp(-128, 127).to(torch.int8)
 ############### YOUR CODE ENDS HERE #################

int8_model_accuracy = evaluate(quantized_model, dataloader['test'],
 extra_preprocess=[extra_preprocess])
print(f"int8 model has accuracy={int8_model_accuracy:.2f}%")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18


VGG(
 (backbone): Sequential(
 (0): QuantizedConv2d()
 (1): QuantizedConv2d()
 (2): QuantizedMaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
 (3): QuantizedConv2d()
 (4): QuantizedConv2d()
 (5): QuantizedMaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
 (6): QuantizedConv2d()
 (7): QuantizedConv2d()
 (8): QuantizedMaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
 (9): QuantizedConv2d()
 (10): QuantizedConv2d()
 (11): QuantizedMaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
 (12): QuantizedAvgPool2d(kernel_size=2, stride=2, padding=0)
 )
 (classifier): QuantizedLinear()
)

1

eval: 0%| | 0/20 [00:00<?, ?it/s]

1

int8 model has accuracy=92.90%

Question 9.2 (5 pts)

Explain why there is no ReLU layer in the linear quantized model.

ReLU layers are fused into the previous convolutional layers during the BatchNorm fusion process (conv-bn fusion). This is a standard preprocessing step in quantization where the BatchNorm parameters are absorbed into the convolutional weights and biases, which inherently applies the activation function effects. Additionally, in the quantized inference framework, ReLU behavior is implicitly handled through the quantization of the output ranges - the quantized values and zero points effectively constrain the output to non-negative values when appropriate, making explicit ReLU modules redundant.

Question 10 (5 pts)

Please compare the advantages and disadvantages of k-means-based quantization and linear quantization. You can discuss from the perspective of accuracy, latency, hardware support, etc.

K-Means Quantization:

Advantages: Better accuracy (clusters weights into discrete values reducing quantization error), can achieve higher compression ratios, works well for weight distribution
Disadvantages: Slower inference (requires codebook lookup for each quantized value), limited hardware support (requires custom implementation), larger memory footprint for storing codebooks

Linear Quantization:

Advantages: Faster inference (simple scale/shift operations), excellent hardware support (native int8 support in most processors), simpler implementation, lower deployment overhead
Disadvantages: Slightly lower accuracy (uniform quantization grid may not match weight distribution well), less flexible for highly non-uniform distributions

Tradeoff: K-means prioritizes accuracy and compression but sacrifices speed and hardware efficiency. Linear quantization prioritizes speed and hardware efficiency with reasonable accuracy. Linear is preferred for deployment; k-means is better for research and extreme compression scenarios.

MIT 6.5940 EfficientML.ai: Lab 1 Pruning

Tue, 24 Mar 2026 00:00:00 +0000

MIT 6.5940 EfficientML.ai: Lab 1 Pruning

This colab notebook provides code and a framework for Lab 1 Pruning. You can work out your solutions here.

Please fill out this feedback form when you finished this lab. We would love to hear your thoughts or feedback on how we can improve this lab!

Goals

In this assignment, you will practice pruning a classical neural network model to reduce both model size and latency. The goals of this assignment are as follows:

Understand the basic concept of pruning
Implement and apply fine-grained pruning
Implement and apply channel pruning
Get a basic understanding of performance improvement (such as speedup) from pruning
Understand the differences and tradeoffs between these pruning approaches

There are two main sections in this lab: Fine-grained Pruning and Channel Pruning.

There are 9 questions in total:

For Fine-grained Pruning, there are 5 questions (Question 1-5).
For Channel Pruning, there are 3 questions (Question 6-8).
Question 9 compares fine-grained pruning and channel pruning.

Setup

First, install the required packages and download the datasets and pretrained model. Here we use CIFAR10 dataset and VGG network which is the same as what we used in the Lab 0 tutorial.

1
2
3


print('Installing torchprofile...')
!pip install torchprofile 1>/dev/null
print('All required packages have been successfully installed!')

Installing torchprofile...
All required packages have been successfully installed!

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24


import copy
import math
import random
import time
from collections import OrderedDict, defaultdict
from typing import Union, List

import numpy as np
import torch
from matplotlib import pyplot as plt
from torch import nn
from torch.optim import *
from torch.optim.lr_scheduler import *
from torch.utils.data import DataLoader
from torchprofile import profile_macs
from torchvision.datasets import *
from torchvision.transforms import *
from tqdm.auto import tqdm

from torchprofile import profile_macs

assert torch.cuda.is_available(), \
"The current runtime does not have CUDA support." \
"Please go to menu bar (Runtime - Change runtime type) and select GPU"

1
2
3


random.seed(0)
np.random.seed(0)
torch.manual_seed(0)

<torch._C.Generator at 0x7a01f05a4a10>

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20


def download_url(url, model_dir='.', overwrite=False):
 import os, sys, ssl
 from urllib.request import urlretrieve
 ssl._create_default_https_context = ssl._create_unverified_context
 target_dir = url.split('/')[-1]
 model_dir = os.path.expanduser(model_dir)
 try:
 if not os.path.exists(model_dir):
 os.makedirs(model_dir)
 model_dir = os.path.join(model_dir, target_dir)
 cached_file = model_dir
 if not os.path.exists(cached_file) or overwrite:
 sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file))
 urlretrieve(url, cached_file)
 return cached_file
 except Exception as e:
 # remove lock file so download can be executed next time.
 os.remove(os.path.join(model_dir, 'download.lock'))
 sys.stderr.write('Failed to download from url %s' % url + '\n' + str(e) + '\n')
 return None

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38


class VGG(nn.Module):
 ARCH = [64, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M']

 def __init__(self) -> None:
 super().__init__()

 layers = []
 counts = defaultdict(int)

 def add(name: str, layer: nn.Module) -> None:
 layers.append((f"{name}{counts[name]}", layer))
 counts[name] += 1

 in_channels = 3
 for x in self.ARCH:
 if x != 'M':
 # conv-bn-relu
 add("conv", nn.Conv2d(in_channels, x, 3, padding=1, bias=False))
 add("bn", nn.BatchNorm2d(x))
 add("relu", nn.ReLU(True))
 in_channels = x
 else:
 # maxpool
 add("pool", nn.MaxPool2d(2))

 self.backbone = nn.Sequential(OrderedDict(layers))
 self.classifier = nn.Linear(512, 10)

 def forward(self, x: torch.Tensor) -> torch.Tensor:
 # backbone: [N, 3, 32, 32] => [N, 512, 2, 2]
 x = self.backbone(x)

 # avgpool: [N, 512, 2, 2] => [N, 512]
 x = x.mean([2, 3])

 # classifier: [N, 512] => [N, 10]
 x = self.classifier(x)
 return x

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32


def train(
 model: nn.Module,
 dataloader: DataLoader,
 criterion: nn.Module,
 optimizer: Optimizer,
 scheduler: LambdaLR,
 callbacks = None
) -> None:
 model.train()

 for inputs, targets in tqdm(dataloader, desc='train', leave=False):
 # Move the data from CPU to GPU
 inputs = inputs.cuda()
 targets = targets.cuda()

 # Reset the gradients (from the last iteration)
 optimizer.zero_grad()

 # Forward inference
 outputs = model(inputs)
 loss = criterion(outputs, targets)

 # Backward propagation
 loss.backward()

 # Update optimizer and LR scheduler
 optimizer.step()
 scheduler.step()

 if callbacks is not None:
 for callback in callbacks:
 callback()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28


@torch.inference_mode()
def evaluate(
 model: nn.Module,
 dataloader: DataLoader,
 verbose=True,
) -> float:
 model.eval()

 num_samples = 0
 num_correct = 0

 for inputs, targets in tqdm(dataloader, desc="eval", leave=False,
 disable=not verbose):
 # Move the data from CPU to GPU
 inputs = inputs.cuda()
 targets = targets.cuda()

 # Inference
 outputs = model(inputs)

 # Convert logits to class indices
 outputs = outputs.argmax(dim=1)

 # Update metrics
 num_samples += targets.size(0)
 num_correct += (outputs == targets).sum()

 return (num_correct / num_samples * 100).item()

Helper Functions (Flops, Model Size calculation, etc.)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49


def get_model_macs(model, inputs) -> int:
 return profile_macs(model, inputs)


def get_sparsity(tensor: torch.Tensor) -> float:
 """
 calculate the sparsity of the given tensor
 sparsity = #zeros / #elements = 1 - #nonzeros / #elements
 """
 return 1 - float(tensor.count_nonzero()) / tensor.numel()


def get_model_sparsity(model: nn.Module) -> float:
 """
 calculate the sparsity of the given model
 sparsity = #zeros / #elements = 1 - #nonzeros / #elements
 """
 num_nonzeros, num_elements = 0, 0
 for param in model.parameters():
 num_nonzeros += param.count_nonzero()
 num_elements += param.numel()
 return 1 - float(num_nonzeros) / num_elements

def get_num_parameters(model: nn.Module, count_nonzero_only=False) -> int:
 """
 calculate the total number of parameters of model
 :param count_nonzero_only: only count nonzero weights
 """
 num_counted_elements = 0
 for param in model.parameters():
 if count_nonzero_only:
 num_counted_elements += param.count_nonzero()
 else:
 num_counted_elements += param.numel()
 return num_counted_elements


def get_model_size(model: nn.Module, data_width=32, count_nonzero_only=False) -> int:
 """
 calculate the model size in bits
 :param data_width: #bits per element
 :param count_nonzero_only: only count nonzero weights
 """
 return get_num_parameters(model, count_nonzero_only) * data_width

Byte = 8
KiB = 1024 * Byte
MiB = 1024 * KiB
GiB = 1024 * MiB

Define misc functions for verification.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52


def test_fine_grained_prune(
 test_tensor=torch.tensor([[-0.46, -0.40, 0.39, 0.19, 0.37],
 [0.00, 0.40, 0.17, -0.15, 0.16],
 [-0.20, -0.23, 0.36, 0.25, 0.03],
 [0.24, 0.41, 0.07, 0.13, -0.15],
 [0.48, -0.09, -0.36, 0.12, 0.45]]),
 test_mask=torch.tensor([[True, True, False, False, False],
 [False, True, False, False, False],
 [False, False, False, False, False],
 [False, True, False, False, False],
 [True, False, False, False, True]]),
 target_sparsity=0.75, target_nonzeros=None):
 def plot_matrix(tensor, ax, title):
 ax.imshow(tensor.cpu().numpy() == 0, vmin=0, vmax=1, cmap='tab20c')
 ax.set_title(title)
 ax.set_yticklabels([])
 ax.set_xticklabels([])
 for i in range(tensor.shape[1]):
 for j in range(tensor.shape[0]):
 text = ax.text(j, i, f'{tensor[i, j].item():.2f}',
 ha="center", va="center", color="k")

 test_tensor = test_tensor.clone()
 fig, axes = plt.subplots(1,2, figsize=(6, 10))
 ax_left, ax_right = axes.ravel()
 plot_matrix(test_tensor, ax_left, 'dense tensor')

 sparsity_before_pruning = get_sparsity(test_tensor)
 mask = fine_grained_prune(test_tensor, target_sparsity)
 sparsity_after_pruning = get_sparsity(test_tensor)
 sparsity_of_mask = get_sparsity(mask)

 plot_matrix(test_tensor, ax_right, 'sparse tensor')
 fig.tight_layout()
 plt.show()

 print('* Test fine_grained_prune()')
 print(f' target sparsity: {target_sparsity:.2f}')
 print(f' sparsity before pruning: {sparsity_before_pruning:.2f}')
 print(f' sparsity after pruning: {sparsity_after_pruning:.2f}')
 print(f' sparsity of pruning mask: {sparsity_of_mask:.2f}')

 if target_nonzeros is None:
 if test_mask.equal(mask):
 print('* Test passed.')
 else:
 print('* Test failed.')
 else:
 if mask.count_nonzero() == target_nonzeros:
 print('* Test passed.')
 else:
 print('* Test failed.')

Load the pretrained model and the CIFAR-10 dataset.

1
2
3
4
5
6


checkpoint_url = "https://hanlab18.mit.edu/files/course/labs/vgg.cifar.pretrained.pth"
checkpoint = torch.load(download_url(checkpoint_url), map_location="cpu")
model = VGG().cuda()
print(f"=> loading checkpoint '{checkpoint_url}'")
model.load_state_dict(checkpoint['state_dict'])
recover_model = lambda: model.load_state_dict(checkpoint['state_dict'])

Downloading: "https://hanlab18.mit.edu/files/course/labs/vgg.cifar.pretrained.pth" to ./vgg.cifar.pretrained.pth


=> loading checkpoint 'https://hanlab18.mit.edu/files/course/labs/vgg.cifar.pretrained.pth'

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26


image_size = 32
transforms = {
 "train": Compose([
 RandomCrop(image_size, padding=4),
 RandomHorizontalFlip(),
 ToTensor(),
 ]),
 "test": ToTensor(),
}
dataset = {}
for split in ["train", "test"]:
 dataset[split] = CIFAR10(
 root="data/cifar10",
 train=(split == "train"),
 download=True,
 transform=transforms[split],
 )
dataloader = {}
for split in ['train', 'test']:
 dataloader[split] = DataLoader(
 dataset[split],
 batch_size=512,
 shuffle=(split == 'train'),
 num_workers=0,
 pin_memory=True,
 )

100%|██████████| 170M/170M [00:04<00:00, 35.0MB/s]

Let’s First Evaluate the Accuracy and Model Size of Dense Model

Neural networks have become ubiquitous in many applications. Here we have loaded a pretrained VGG model for classifying images in CIFAR10 dataset.

Let’s first evaluate the accuracy and model size of this model.

1
2
3
4


dense_model_accuracy = evaluate(model, dataloader['test'])
dense_model_size = get_model_size(model)
print(f"dense model has accuracy={dense_model_accuracy:.2f}%")
print(f"dense model has size={dense_model_size/MiB:.2f} MiB")

eval: 0%| | 0/20 [00:00<?, ?it/s]


dense model has accuracy=92.95%
dense model has size=35.20 MiB

While large neural networks are very powerful, their size consumes considerable storage, memory bandwidth, and computational resources. As we can see from the results above, a model for the task as simple as classifying $32\times32$ images into 10 classes can be as large as 35 MiB. For embedded mobile applications, these resource demands become prohibitive.

Therefore, neural network pruning is exploited to facilitates storage and transmission of mobile applications incorporating DNNs.

The goal of pruning is to reduce the model size while maintaining the accuracy.

Let’s see the distribution of weight values

Before we jump into pruning, let’s see the distribution of weight values in the dense model.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24


def plot_weight_distribution(model, bins=256, count_nonzero_only=False):
 fig, axes = plt.subplots(3,3, figsize=(10, 6))
 axes = axes.ravel()
 plot_index = 0
 for name, param in model.named_parameters():
 if param.dim() > 1:
 ax = axes[plot_index]
 if count_nonzero_only:
 param_cpu = param.detach().view(-1).cpu()
 param_cpu = param_cpu[param_cpu != 0].view(-1)
 ax.hist(param_cpu, bins=bins, density=True,
 color = 'blue', alpha = 0.5)
 else:
 ax.hist(param.detach().view(-1).cpu(), bins=bins, density=True,
 color = 'blue', alpha = 0.5)
 ax.set_xlabel(name)
 ax.set_ylabel('density')
 plot_index += 1
 fig.suptitle('Histogram of Weights')
 fig.tight_layout()
 fig.subplots_adjust(top=0.925)
 plt.show()

plot_weight_distribution(model)

Question 1 (10 pts)

Please answer the following questions using the information in the above histograms of weights.

Question 1.1 (5 pts)

What are the common characteristics of the weight distribution in the different layers?

My Answer: The weight distributions of different layers are usually centered around 0, roughly symmetric, and have a sharp peak with long tails: many weights are concentrated near 0, while only a small number have large magnitudes.

Question 1.2 (5 pts)

How do these characteristics help pruning?

My Answer: Since most parameters have small absolute values (close to 0), pruning these small-magnitude parameters usually has limited impact on outputs, so we can significantly increase sparsity and compression with only minor accuracy loss.

Fine-grained Pruning

In this section, we will implement and perform fine-grained pruning.

Fine-grained pruning removes the synapses with lowest importance. The weight tensor $W$ will become sparse after fine-grained pruning, which can be described with sparsity:

$\mathrm{sparsity} := \#\mathrm{Zeros} / \#W = 1 - \#\mathrm{Nonzeros} / \#W$

where $\#W$ is the number of elements in $W$.

In practice, given the target sparsity $s$, the weight tensor $W$ is multiplied with a binary mask $M$ to disregard removed weight:

$v_{\mathrm{thr}} = \texttt{kthvalue}(Importance, \#W \cdot s)$

$M = Importance > v_{\mathrm{thr}}$

$W = W \cdot M$

where $Importance$ is importance tensor with the same shape of $W$, $\texttt{kthvalue}(X, k)$ finds the $k$-th smallest value of tensor $X$, $v_{\mathrm{thr}}$ is the threshold value.

Magnitude-based Pruning

For fine-grained pruning, a widely-used importance is the magnitude of weight value, i.e.,

$Importance=|W|$

This is known as Magnitude-based Pruning (see Learning both Weights and Connections for Efficient Neural Networks).

Question 2 (15 pts)

Please complete the following magnitude-based fine-grained pruning function.

Hint:

In step 1, we calculate the number of zeros (num_zeros) after pruning. Note that num_zeros should be an integer. You could use either round() or int() to convert a floating number into an integer. Here we use round().
In step 2, we calculate the importance of weight tensor. Pytorch provides torch.abs(), torch.Tensor.abs(), torch.Tensor.abs_() APIs.
In step 3, we calculate the pruning threshold so that all synapses with importance smaller than threshold will be removed. Pytorch provides torch.kthvalue(), torch.Tensor.kthvalue(), torch.topk() APIs.
In step 4, we calculate the pruning mask based on the threshold. 1 in the mask indicates the synapse will be kept, and 0 in the mask indicates the synapse will be removed. mask = importance > threshold. Pytorch provides torch.gt() API.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33


def fine_grained_prune(tensor: torch.Tensor, sparsity : float) -> torch.Tensor:
 """
 magnitude-based pruning for single tensor
 :param tensor: torch.(cuda.)Tensor, weight of conv/fc layer
 :param sparsity: float, pruning sparsity
 sparsity = #zeros / #elements = 1 - #nonzeros / #elements
 :return:
 torch.(cuda.)Tensor, mask for zeros
 """
 sparsity = min(max(0.0, sparsity), 1.0)
 if sparsity == 1.0:
 tensor.zero_()
 return torch.zeros_like(tensor)
 elif sparsity == 0.0:
 return torch.ones_like(tensor)

 num_elements = tensor.numel()

 ##################### MY CODE STARTS HERE #####################
 # Step 1: calculate the #zeros (please use round())
 num_zeros = round(num_elements * sparsity)
 # Step 2: calculate the importance of weight
 importance = tensor.abs()
 # Step 3: calculate the pruning threshold
 threshold = torch.kthvalue(importance.view(-1), num_zeros).values
 # Step 4: get binary mask (1 for nonzeros, 0 for zeros)
 mask = importance > threshold
 ##################### MY CODE ENDS HERE #######################

 # Step 5: apply mask to prune the tensor
 tensor.mul_(mask)

 return mask

Let’s verify the functionality of defined fine-grained pruning by applying the function above on a dummy tensor.

1

test_fine_grained_prune()

* Test fine_grained_prune()
 target sparsity: 0.75
 sparsity before pruning: 0.04
 sparsity after pruning: 0.76
 sparsity of pruning mask: 0.76
* Test passed.

Question 3 (5 pts)

The last cell plots the tensor before and after pruning. Nonzeros are rendered in blue while zeros are rendered in gray. Please modify the value of target_sparsity in the following code cell so that there are only 10 nonzeros in the sparse tensor after pruning.

1
2
3
4


##################### MY CODE STARTS HERE #####################
target_sparsity = 0.6 # 25 * (1 - 0.6) = 10 nonzeros
##################### MY CODE ENDS HERE #####################
test_fine_grained_prune(target_sparsity=target_sparsity, target_nonzeros=10)

* Test fine_grained_prune()
 target sparsity: 0.60
 sparsity before pruning: 0.04
 sparsity after pruning: 0.60
 sparsity of pruning mask: 0.60
* Test passed.

We now wrap the fine-grained pruning function into a class for pruning the whole model. In class FineGrainedPruner, we have to keep a record of the pruning masks so that we could apply the masks whenever the model weights change to make sure the model keep sparse all the time.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18


class FineGrainedPruner:
 def __init__(self, model, sparsity_dict):
 self.masks = FineGrainedPruner.prune(model, sparsity_dict)

 @torch.no_grad()
 def apply(self, model):
 for name, param in model.named_parameters():
 if name in self.masks:
 param *= self.masks[name]

 @staticmethod
 @torch.no_grad()
 def prune(model, sparsity_dict):
 masks = dict()
 for name, param in model.named_parameters():
 if param.dim() > 1: # we only prune conv and fc weights
 masks[name] = fine_grained_prune(param, sparsity_dict[name])
 return masks

Sensitivity Scan

Different layers contribute differently to the model performance. It is challenging to decide the proper sparsity for each layer. A widely used approach is sensitivity scan.

During the sensitivity scan, at each time, we will only prune one layer to see the accuracy degradation. By scanning different sparsities, we could draw the sensitivity curve (i.e., accuracy vs. sparsity) of the corresponding layer.

Here is an example figure for sensitivity curves. The x-axis is the sparsity or the percentage of #parameters dropped (i.e., sparsity). The y-axis is the validation accuracy. (This is Figure 6 in Learning both Weights and Connections for Efficient Neural Networks)

The following code cell defines the sensitivity scan function that returns the sparsities scanned, and a list of accuracies for each weight to be pruned.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21


@torch.no_grad()
def sensitivity_scan(model, dataloader, scan_step=0.1, scan_start=0.4, scan_end=1.0, verbose=True):
 sparsities = np.arange(start=scan_start, stop=scan_end, step=scan_step)
 accuracies = []
 named_conv_weights = [(name, param) for (name, param) \
 in model.named_parameters() if param.dim() > 1]
 for i_layer, (name, param) in enumerate(named_conv_weights):
 param_clone = param.detach().clone()
 accuracy = []
 for sparsity in tqdm(sparsities, desc=f'scanning {i_layer}/{len(named_conv_weights)} weight - {name}'):
 fine_grained_prune(param.detach(), sparsity=sparsity)
 acc = evaluate(model, dataloader, verbose=False)
 if verbose:
 print(f'\r sparsity={sparsity:.2f}: accuracy={acc:.2f}%', end='')
 # restore
 param.copy_(param_clone)
 accuracy.append(acc)
 if verbose:
 print(f'\r sparsity=[{",".join(["{:.2f}".format(x) for x in sparsities])}]: accuracy=[{", ".join(["{:.2f}%".format(x) for x in accuracy])}]', end='')
 accuracies.append(accuracy)
 return sparsities, accuracies

Please run the following cells to plot the sensitivity curves. It should take around 2 minutes to finish.

1
2


sparsities, accuracies = sensitivity_scan(
 model, dataloader['test'], scan_step=0.1, scan_start=0.4, scan_end=1.0)

scanning 0/9 weight - backbone.conv0.weight: 0%| | 0/6 [00:00<?, ?it/s]


 sparsity=[0.40,0.50,0.60,0.70,0.80,0.90]: accuracy=[92.42%, 91.19%, 87.55%, 83.39%, 69.41%, 31.81%]


scanning 1/9 weight - backbone.conv1.weight: 0%| | 0/6 [00:00<?, ?it/s]


 sparsity=[0.40,0.50,0.60,0.70,0.80,0.90]: accuracy=[92.93%, 92.88%, 92.71%, 92.40%, 91.32%, 84.78%]


scanning 2/9 weight - backbone.conv2.weight: 0%| | 0/6 [00:00<?, ?it/s]


 sparsity=[0.40,0.50,0.60,0.70,0.80,0.90]: accuracy=[92.94%, 92.64%, 92.46%, 91.77%, 89.85%, 78.56%]


scanning 3/9 weight - backbone.conv3.weight: 0%| | 0/6 [00:00<?, ?it/s]


 sparsity=[0.40,0.50,0.60,0.70,0.80,0.90]: accuracy=[92.86%, 92.72%, 92.23%, 91.09%, 85.35%, 51.31%]


scanning 4/9 weight - backbone.conv4.weight: 0%| | 0/6 [00:00<?, ?it/s]


 sparsity=[0.40,0.50,0.60,0.70,0.80,0.90]: accuracy=[92.88%, 92.68%, 92.22%, 89.47%, 76.86%, 38.78%]


scanning 5/9 weight - backbone.conv5.weight: 0%| | 0/6 [00:00<?, ?it/s]


 sparsity=[0.40,0.50,0.60,0.70,0.80,0.90]: accuracy=[92.92%, 92.71%, 92.63%, 91.88%, 89.90%, 82.19%]


scanning 6/9 weight - backbone.conv6.weight: 0%| | 0/6 [00:00<?, ?it/s]


 sparsity=[0.40,0.50,0.60,0.70,0.80,0.90]: accuracy=[92.94%, 92.86%, 92.65%, 92.10%, 90.58%, 83.65%]


scanning 7/9 weight - backbone.conv7.weight: 0%| | 0/6 [00:00<?, ?it/s]


 sparsity=[0.40,0.50,0.60,0.70,0.80,0.90]: accuracy=[92.94%, 92.92%, 92.88%, 92.81%, 92.63%, 91.34%]


scanning 8/9 weight - classifier.weight: 0%| | 0/6 [00:00<?, ?it/s]


 sparsity=[0.40,0.50,0.60,0.70,0.80,0.90]: accuracy=[92.91%, 92.83%, 92.81%, 92.97%, 92.68%, 92.52%]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27


def plot_sensitivity_scan(sparsities, accuracies, dense_model_accuracy):
 lower_bound_accuracy = 100 - (100 - dense_model_accuracy) * 1.5
 fig, axes = plt.subplots(3, int(math.ceil(len(accuracies) / 3)),figsize=(15,8))
 axes = axes.ravel()
 plot_index = 0
 for name, param in model.named_parameters():
 if param.dim() > 1:
 ax = axes[plot_index]
 curve = ax.plot(sparsities, accuracies[plot_index])
 line = ax.plot(sparsities, [lower_bound_accuracy] * len(sparsities))
 ax.set_xticks(np.arange(start=0.4, stop=1.0, step=0.1))
 ax.set_ylim(80, 95)
 ax.set_title(name)
 ax.set_xlabel('sparsity')
 ax.set_ylabel('top-1 accuracy')
 ax.legend([
 'accuracy after pruning',
 f'{lower_bound_accuracy / dense_model_accuracy * 100:.0f}% of dense model accuracy'
 ])
 ax.grid(axis='x')
 plot_index += 1
 fig.suptitle('Sensitivity Curves: Validation Accuracy vs. Pruning Sparsity')
 fig.tight_layout()
 fig.subplots_adjust(top=0.925)
 plt.show()

plot_sensitivity_scan(sparsities, accuracies, dense_model_accuracy)

Question 4 (15 pts)

Please answer the following questions using the information in the above sensitivity curves.

Question 4.1 (5 pts)

What’s the relationship between pruning sparsity and model accuracy? (i.e., does accuracy increase or decrease when sparsity becomes higher?)

My Answer: In general, higher pruning sparsity leads to lower accuracy; in the low-to-medium sparsity range, the drop is gradual, but after a certain turning point, accuracy degrades much faster.

Question 4.2 (5 pts)

Do all the layers have the same sensitivity?

My Answer: No. Different layers have very different sensitivity to sparsity; layers closer to the output or more critical layers are usually more sensitive.

Question 4.3 (5 pts)

Which layer is the most sensitive to the pruning sparsity?

My Answer: The most sensitive layers are typically high-semantic layers near the back end (for example, backbone.conv6/backbone.conv7 or the classifier), because they directly affect the final decision boundary.

#Parameters of each layer

In addition to accuracy, the number of each layer’s parameters also affects the decision on sparsity selection. Layers with more #parameters require larger sparsities.

Please run the following code cell to plot the distribution of #parameters in the whole model.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15


def plot_num_parameters_distribution(model):
 num_parameters = dict()
 for name, param in model.named_parameters():
 if param.dim() > 1:
 num_parameters[name] = param.numel()
 fig = plt.figure(figsize=(8, 6))
 plt.grid(axis='y')
 plt.bar(list(num_parameters.keys()), list(num_parameters.values()))
 plt.title('#Parameter Distribution')
 plt.ylabel('Number of Parameters')
 plt.xticks(rotation=60)
 plt.tight_layout()
 plt.show()

plot_num_parameters_distribution(model)

Select Sparsity Based on Sensitivity Curves and #Parameters Distribution

Question 5 (10 pts)

Based on the sensitivity curves and the distribution of #parameters in the model, please select the sparsity for each layer.

Note that the overall compression ratio of pruned model mostly depends on the layers with larger #parameters, and different layers have different sensitivity to pruning (see Question 4).

Please make sure that after pruning, the sparse model is 25% of the size of the dense model, and validation accuracy is higher than 92.5 after finetuning.

Hint:

The layer with more #parameters should have larger sparsity. (see Figure #Parameter Distribution)
The layer that is sensitive to the pruning sparsity (i.e., the accuracy will drop quickly as sparsity becomes higher) should have smaller sparsity. (see Figure Sensitivity Curves)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17


recover_model()

sparsity_dict = {
##################### MY CODE STARTS HERE #####################
 # please modify the sparsity value of each layer
 # please DO NOT modify the key of sparsity_dict
 'backbone.conv0.weight': 0.20,
 'backbone.conv1.weight': 0.30,
 'backbone.conv2.weight': 0.50,
 'backbone.conv3.weight': 0.60,
 'backbone.conv4.weight': 0.70,
 'backbone.conv5.weight': 0.75,
 'backbone.conv6.weight': 0.80,
 'backbone.conv7.weight': 0.80,
 'classifier.weight': 0.50
##################### MY CODE ENDS HERE #######################
}

Please run the following cell to prune the model according to your defined sparsity_dict, and print the information of sparse model.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15


pruner = FineGrainedPruner(model, sparsity_dict)
print(f'After pruning with sparsity dictionary')
for name, sparsity in sparsity_dict.items():
 print(f' {name}: {sparsity:.2f}')
print(f'The sparsity of each layer becomes')
for name, param in model.named_parameters():
 if name in sparsity_dict:
 print(f' {name}: {get_sparsity(param):.2f}')

sparse_model_size = get_model_size(model, count_nonzero_only=True)
print(f"Sparse model has size={sparse_model_size / MiB:.2f} MiB = {sparse_model_size / dense_model_size * 100:.2f}% of dense model size")
sparse_model_accuracy = evaluate(model, dataloader['test'])
print(f"Sparse model has accuracy={sparse_model_accuracy:.2f}% before fintuning")

plot_weight_distribution(model, count_nonzero_only=True)

After pruning with sparsity dictionary
 backbone.conv0.weight: 0.20
 backbone.conv1.weight: 0.30
 backbone.conv2.weight: 0.50
 backbone.conv3.weight: 0.60
 backbone.conv4.weight: 0.70
 backbone.conv5.weight: 0.75
 backbone.conv6.weight: 0.80
 backbone.conv7.weight: 0.80
 classifier.weight: 0.50
The sparsity of each layer becomes
 backbone.conv0.weight: 0.20
 backbone.conv1.weight: 0.30
 backbone.conv2.weight: 0.50
 backbone.conv3.weight: 0.60
 backbone.conv4.weight: 0.70
 backbone.conv5.weight: 0.75
 backbone.conv6.weight: 0.80
 backbone.conv7.weight: 0.80
 classifier.weight: 0.50
Sparse model has size=8.90 MiB = 25.27% of dense model size



eval: 0%| | 0/20 [00:00<?, ?it/s]


Sparse model has accuracy=88.90% before fintuning

Finetune the fine-grained pruned model

As we can see from the outputs of previous cell, even though fine-grained pruning reduces the most of model weights, the accuracy of model also dropped. Therefore, we have to finetune the sparse model to recover the accuracy.

Please run the following cell to finetune the sparse model. It should take around 3 minutes to finish.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19


num_finetune_epochs = 5
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, num_finetune_epochs)
criterion = nn.CrossEntropyLoss()

best_sparse_model_checkpoint = dict()
best_accuracy = 0
print(f'Finetuning Fine-grained Pruned Sparse Model')
for epoch in range(num_finetune_epochs):
 # At the end of each train iteration, we have to apply the pruning mask
 # to keep the model sparse during the training
 train(model, dataloader['train'], criterion, optimizer, scheduler,
 callbacks=[lambda: pruner.apply(model)])
 accuracy = evaluate(model, dataloader['test'])
 is_best = accuracy > best_accuracy
 if is_best:
 best_sparse_model_checkpoint['state_dict'] = copy.deepcopy(model.state_dict())
 best_accuracy = accuracy
 print(f' Epoch {epoch+1} Accuracy {accuracy:.2f}% / Best Accuracy: {best_accuracy:.2f}%')

Finetuning Fine-grained Pruned Sparse Model



train: 0%| | 0/98 [00:00<?, ?it/s]



eval: 0%| | 0/20 [00:00<?, ?it/s]


 Epoch 1 Accuracy 92.75% / Best Accuracy: 92.75%



train: 0%| | 0/98 [00:00<?, ?it/s]



eval: 0%| | 0/20 [00:00<?, ?it/s]


 Epoch 2 Accuracy 92.86% / Best Accuracy: 92.86%



train: 0%| | 0/98 [00:00<?, ?it/s]



eval: 0%| | 0/20 [00:00<?, ?it/s]


 Epoch 3 Accuracy 92.78% / Best Accuracy: 92.86%



train: 0%| | 0/98 [00:00<?, ?it/s]



eval: 0%| | 0/20 [00:00<?, ?it/s]


 Epoch 4 Accuracy 92.83% / Best Accuracy: 92.86%



train: 0%| | 0/98 [00:00<?, ?it/s]



eval: 0%| | 0/20 [00:00<?, ?it/s]


 Epoch 5 Accuracy 92.84% / Best Accuracy: 92.86%

Run the following cell to see the information of best finetuned sparse model.

1
2
3
4
5
6


# load the best sparse model checkpoint to evaluate the final performance
model.load_state_dict(best_sparse_model_checkpoint['state_dict'])
sparse_model_size = get_model_size(model, count_nonzero_only=True)
print(f"Sparse model has size={sparse_model_size / MiB:.2f} MiB = {sparse_model_size / dense_model_size * 100:.2f}% of dense model size")
sparse_model_accuracy = evaluate(model, dataloader['test'])
print(f"Sparse model has accuracy={sparse_model_accuracy:.2f}% after fintuning")

Sparse model has size=8.90 MiB = 25.27% of dense model size



eval: 0%| | 0/20 [00:00<?, ?it/s]


Sparse model has accuracy=92.86% after fintuning

Channel Pruning

In this section, we will implement the channel pruning. Channel pruning removes an entire channel, so that it can achieve inference speed up on existing hardware like GPUs. Similarly, we remove the channels whose weights are of smaller magnitudes (measured by Frobenius norm).

1
2
3
4
5


# firstly, let's restore the model weights to the original dense version
# and check the validation accuracy
recover_model()
dense_model_accuracy = evaluate(model, dataloader['test'])
print(f"dense model has accuracy={dense_model_accuracy:.2f}%")

eval: 0%| | 0/20 [00:00<?, ?it/s]


dense model has accuracy=92.95%

Remove Channel Weights

Unlike fine-grained pruning, we can remove the weights entirely from the tensor in channel pruning. That is to say, the number of output channels is reduced:

$\#\mathrm{out\_channels}_{\mathrm{new}} = \#\mathrm{out\_channels}_{\mathrm{origin}} \cdot (1 - \mathrm{sparsity})$

The weight tensor $W$ is still dense after channel pruning. Thus, we will refer to sparsity as prune ratio.

Like fine-grained pruning, we can use different pruning rates for different layers. However, we use a uniform pruning rate for all the layers for now. We are targeting 2x computation reduction, which is roughly 30% uniform pruning rate (think about why).

Feel free to try out different pruning ratios per layer at the end of this section. You can pass in a list of ratios to the channel_prune function.

Question 6 (10 pts)

Please complete the following functions for channel pruning.

Here we naively prune all output channels other than the first $\#\mathrm{out\_channels}_{\mathrm{new}}$ channels.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53


def get_num_channels_to_keep(channels: int, prune_ratio: float) -> int:
 """A function to calculate the number of layers to PRESERVE after pruning
 Note that preserve_rate = 1. - prune_ratio
 """
 ##################### MY CODE STARTS HERE #####################
 return int(round(channels * (1.0 - prune_ratio)))
 ##################### MY CODE ENDS HERE #####################

@torch.no_grad()
def channel_prune(model: nn.Module,
 prune_ratio: Union[List, float]) -> nn.Module:
 """Apply channel pruning to each of the conv layer in the backbone
 Note that for prune_ratio, we can either provide a floating-point number,
 indicating that we use a uniform pruning rate for all layers, or a list of
 numbers to indicate per-layer pruning rate.
 """
 # sanity check of provided prune_ratio
 assert isinstance(prune_ratio, (float, list))
 n_conv = len([m for m in model.backbone if isinstance(m, nn.Conv2d)])
 # note that for the ratios, it affects the previous conv output and next
 # conv input, i.e., conv0 - ratio0 - conv1 - ratio1-...
 if isinstance(prune_ratio, list):
 assert len(prune_ratio) == n_conv - 1
 else: # convert float to list
 prune_ratio = [prune_ratio] * (n_conv - 1)

 # we prune the convs in the backbone with a uniform ratio
 model = copy.deepcopy(model) # prevent overwrite
 # we only apply pruning to the backbone features
 all_convs = [m for m in model.backbone if isinstance(m, nn.Conv2d)]
 all_bns = [m for m in model.backbone if isinstance(m, nn.BatchNorm2d)]
 # apply pruning. we naively keep the first k channels
 assert len(all_convs) == len(all_bns)
 for i_ratio, p_ratio in enumerate(prune_ratio):
 prev_conv = all_convs[i_ratio]
 prev_bn = all_bns[i_ratio]
 next_conv = all_convs[i_ratio + 1]
 original_channels = prev_conv.out_channels # same as next_conv.in_channels
 n_keep = get_num_channels_to_keep(original_channels, p_ratio)

 # prune the output of the previous conv and bn
 prev_conv.weight.set_(prev_conv.weight.detach()[:n_keep])
 prev_bn.weight.set_(prev_bn.weight.detach()[:n_keep])
 prev_bn.bias.set_(prev_bn.bias.detach()[:n_keep])
 prev_bn.running_mean.set_(prev_bn.running_mean.detach()[:n_keep])
 prev_bn.running_var.set_(prev_bn.running_var.detach()[:n_keep])

 # prune the input of the next conv (hint: just one line of code)
 ##################### YOUR CODE STARTS HERE #####################
 next_conv.weight.set_(next_conv.weight.detach()[:, :n_keep])
 ##################### YOUR CODE ENDS HERE #####################

 return model

Run the following cell to perform a sanity check to make sure the implementation is correct.

1
2
3
4
5


dummy_input = torch.randn(1, 3, 32, 32).cuda()
pruned_model = channel_prune(model, prune_ratio=0.3)
pruned_macs = get_model_macs(pruned_model, dummy_input)
assert pruned_macs == 305388064
print('* Check passed. Right MACs for the pruned model.')

* Check passed. Right MACs for the pruned model.

Now let’s evaluate the performance of the model after uniform channel pruning with 30% pruning rate.

As you may see, directly removing 30% of the channels leads to low accuracy.

1
2


pruned_model_accuracy = evaluate(pruned_model, dataloader['test'])
print(f"pruned model has accuracy={pruned_model_accuracy:.2f}%")

eval: 0%| | 0/20 [00:00<?, ?it/s]


pruned model has accuracy=28.14%

Ranking Channels by Importance

As you can see, removing the first 30% of channels in all layers leads to significant accuracy reduction. One potential method to remedy the issue is to find the less important channel weights to remove. A popular criterion for importance is to use the Frobenius norm of the weights corresponding to each input channel:

$importance_{i} = \|W_{i}\|_2, \;\; i = 0, 1, 2,\cdots, \#\mathrm{in\_channels}-1$

We can sort the channel weights from more important to less important, and then keep the frst $k$ channels for each layer.

Question 7 (15 pts)

Please complete the following functions for sorting the weight tensor based on the Frobenius norm.

Hint:

To calculate Frobenius norm of a tensor, Pytorch provides torch.norm APIs.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48


# function to sort the channels from important to non-important
def get_input_channel_importance(weight):
 in_channels = weight.shape[1]
 importances = []
 # compute the importance for each input channel
 for i_c in range(weight.shape[1]):
 channel_weight = weight.detach()[:, i_c]
 ##################### MY CODE STARTS HERE #####################
 importance = torch.norm(channel_weight, p=2)
 ##################### MY CODE ENDS HERE #####################
 importances.append(importance.view(1))
 return torch.cat(importances)

@torch.no_grad()
def apply_channel_sorting(model):
 model = copy.deepcopy(model) # do not modify the original model
 # fetch all the conv and bn layers from the backbone
 all_convs = [m for m in model.backbone if isinstance(m, nn.Conv2d)]
 all_bns = [m for m in model.backbone if isinstance(m, nn.BatchNorm2d)]
 # iterate through conv layers
 for i_conv in range(len(all_convs) - 1):
 # each channel sorting index, we need to apply it to:
 # - the output dimension of the previous conv
 # - the previous BN layer
 # - the input dimension of the next conv (we compute importance here)
 prev_conv = all_convs[i_conv]
 prev_bn = all_bns[i_conv]
 next_conv = all_convs[i_conv + 1]
 # note that we always compute the importance according to input channels
 importance = get_input_channel_importance(next_conv.weight)
 # sorting from large to small
 sort_idx = torch.argsort(importance, descending=True)

 # apply to previous conv and its following bn
 prev_conv.weight.copy_(torch.index_select(
 prev_conv.weight.detach(), 0, sort_idx))
 for tensor_name in ['weight', 'bias', 'running_mean', 'running_var']:
 tensor_to_apply = getattr(prev_bn, tensor_name)
 tensor_to_apply.copy_(
 torch.index_select(tensor_to_apply.detach(), 0, sort_idx)
 )

 # apply to the next conv input (hint: one line of code)
 ##################### YOUR CODE STARTS HERE #####################
 next_conv.weight.copy_(torch.index_select(next_conv.weight.detach(), 1, sort_idx))
 ##################### YOUR CODE ENDS HERE #####################

 return model

Now run the following cell to sanity check if the results are correct.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


print('Before sorting...')
dense_model_accuracy = evaluate(model, dataloader['test'])
print(f"dense model has accuracy={dense_model_accuracy:.2f}%")

print('After sorting...')
sorted_model = apply_channel_sorting(model)
sorted_model_accuracy = evaluate(sorted_model, dataloader['test'])
print(f"sorted model has accuracy={sorted_model_accuracy:.2f}%")

# make sure accuracy does not change after sorting, since it is
# equivalent transform
assert abs(sorted_model_accuracy - dense_model_accuracy) < 0.1
print('* Check passed.')

Before sorting...



eval: 0%| | 0/20 [00:00<?, ?it/s]


dense model has accuracy=92.95%
After sorting...



eval: 0%| | 0/20 [00:00<?, ?it/s]


sorted model has accuracy=92.95%
* Check passed.

Finally, we compare the pruned models’ accuracy with and without sorting.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


channel_pruning_ratio = 0.3 # pruned-out ratio

print(" * Without sorting...")
pruned_model = channel_prune(model, channel_pruning_ratio)
pruned_model_accuracy = evaluate(pruned_model, dataloader['test'])
print(f"pruned model has accuracy={pruned_model_accuracy:.2f}%")


print(" * With sorting...")
sorted_model = apply_channel_sorting(model)
pruned_model = channel_prune(sorted_model, channel_pruning_ratio)
pruned_model_accuracy = evaluate(pruned_model, dataloader['test'])
print(f"pruned model has accuracy={pruned_model_accuracy:.2f}%")

 * Without sorting...



eval: 0%| | 0/20 [00:00<?, ?it/s]


pruned model has accuracy=28.14%
 * With sorting...



eval: 0%| | 0/20 [00:00<?, ?it/s]


pruned model has accuracy=36.81%

As you can see, the channel sorting can slightly improve the pruned model’s accuracy, but there is still a huge degrade, which is quite common for channel pruning. But luckily, we can perform fine-tuning to recover the accuracy.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


num_finetune_epochs = 5
optimizer = torch.optim.SGD(pruned_model.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, num_finetune_epochs)
criterion = nn.CrossEntropyLoss()

best_accuracy = 0
for epoch in range(num_finetune_epochs):
 train(pruned_model, dataloader['train'], criterion, optimizer, scheduler)
 accuracy = evaluate(pruned_model, dataloader['test'])
 is_best = accuracy > best_accuracy
 if is_best:
 best_accuracy = accuracy
 print(f'Epoch {epoch+1} Accuracy {accuracy:.2f}% / Best Accuracy: {best_accuracy:.2f}%')

train: 0%| | 0/98 [00:00<?, ?it/s]



eval: 0%| | 0/20 [00:00<?, ?it/s]


Epoch 1 Accuracy 91.65% / Best Accuracy: 91.65%



train: 0%| | 0/98 [00:00<?, ?it/s]



eval: 0%| | 0/20 [00:00<?, ?it/s]


Epoch 2 Accuracy 92.08% / Best Accuracy: 92.08%



train: 0%| | 0/98 [00:00<?, ?it/s]



eval: 0%| | 0/20 [00:00<?, ?it/s]


Epoch 3 Accuracy 92.01% / Best Accuracy: 92.08%



train: 0%| | 0/98 [00:00<?, ?it/s]



eval: 0%| | 0/20 [00:00<?, ?it/s]


Epoch 4 Accuracy 92.14% / Best Accuracy: 92.14%



train: 0%| | 0/98 [00:00<?, ?it/s]



eval: 0%| | 0/20 [00:00<?, ?it/s]


Epoch 5 Accuracy 92.10% / Best Accuracy: 92.14%

Measure acceleration from pruning

After fine-tuning, the model almost recovers the accuracy. You may have already learned that channel pruning is usually more difficult to recover accuracy compared to fine-grained pruning. However, it directly leads to a smaller model size and smaller computation without specialized model format. It can also run faster on GPUs. Now we compare the model size, computation, and latency of the pruned model.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51


# helper functions to measure latency of a regular PyTorch models.
# Unlike fine-grained pruning, channel pruning
# can directly leads to model size reduction and speed up.
@torch.no_grad()
def measure_latency(model, dummy_input, n_warmup=20, n_test=100):
 model.eval()
 # warmup
 for _ in range(n_warmup):
 _ = model(dummy_input)
 # real test
 t1 = time.time()
 for _ in range(n_test):
 _ = model(dummy_input)
 t2 = time.time()
 return (t2 - t1) / n_test # average latency

table_template = "{:<15} {:<15} {:<15} {:<15}"
print (table_template.format('', 'Original','Pruned','Reduction Ratio'))

# 1. measure the latency of the original model and the pruned model on CPU
# which simulates inference on an edge device
dummy_input = torch.randn(1, 3, 32, 32).to('cpu')
pruned_model = pruned_model.to('cpu')
model = model.to('cpu')

pruned_latency = measure_latency(pruned_model, dummy_input)
original_latency = measure_latency(model, dummy_input)
print(table_template.format('Latency (ms)',
 round(original_latency * 1000, 1),
 round(pruned_latency * 1000, 1),
 round(original_latency / pruned_latency, 1)))

# 2. measure the computation (MACs)
original_macs = get_model_macs(model, dummy_input)
pruned_macs = get_model_macs(pruned_model, dummy_input)
print(table_template.format('MACs (M)',
 round(original_macs / 1e6),
 round(pruned_macs / 1e6),
 round(original_macs / pruned_macs, 1)))

# 3. measure the model size (params)
original_param = get_num_parameters(model)
pruned_param = get_num_parameters(pruned_model)
print(table_template.format('Param (M)',
 round(original_param / 1e6, 2),
 round(pruned_param / 1e6, 2),
 round(original_param / pruned_param, 1)))

# put model back to cuda
pruned_model = pruned_model.to('cuda')
model = model.to('cuda')

 Original Pruned Reduction Ratio
Latency (ms) 18.1 10.7 1.7 
MACs (M) 606 305 2.0 
Param (M) 9.23 5.01 1.8

Question 8 (10 pts)

Please answer the following questions using the information in the previous code cell.

Question 8.1 (5 pts)

Explain why removing 30% of channels roughly leads to 50% computation reduction.

My Answer: Convolutional computation is approximately proportional to Cin * Cout * K * K * H * W. Channel pruning reduces both the previous layer’s Cout and the next layer’s Cin, so after removing 30% of channels, the related term becomes about 0.7 * 0.7 = 0.49, meaning computation is roughly halved.

Question 8.2 (5 pts)

Explain why the latency reduction ratio is slightly smaller than computation reduction.

My Answer: Latency is not determined only by MACs; it is also affected by memory access, cache hit rate, kernel launch overhead, hardware parallelism, and framework scheduling overhead. So real latency reduction is often smaller than the theoretical compute reduction.

Compare Fine-grained Pruning and Channel Pruning

Question 9 (10 pts)

After all experiments in this lab, you may have become familiar with both fine-grained pruning and channel pruning.

Please answer the following questions using what you have learned from the lectures and this lab.

Question 9.1 (5 pts)

What are the advantages and disadvantages of fine-grained pruning and channel pruning? You can discuss from the perspective of compression ratio, accuracy, latency, hardware support (i.e., requiring specialized hardware accelerator), etc.

My Answer: Fine-grained pruning often gives higher compression and better accuracy retention, but it produces unstructured sparsity that is hard to accelerate on general hardware. Channel pruning creates structured sparsity, which is easier to accelerate on CPU/GPU/mobile NPU, but at the same compression ratio it usually causes larger accuracy loss.

Question 9.2 (5 pts)

If you want to make your model run faster on a smartphone, which pruning method will you use? Why?

My Answer: Prefer channel pruning. Smartphone hardware and inference frameworks usually support “structurally narrower networks” better, which directly reduces tensor sizes and operator computation, leading to more stable and reproducible on-device latency gains.

B31DD Course Project — Traffic Light System

Sun, 08 Dec 2024 00:00:00 +0000

B31DD Mini-Project Report

Traffic Light Display System with Servo Control and Serial Monitoring

1. Project Introduction

The Traffic Light Display System is an embedded project designed to simulate a real-world traffic light controller using an AVR microcontroller. This system manages traffic flow by controlling red, yellow, and green LEDs, incorporates a servo motor to represent mechanical signaling (e.g., pedestrian crossing), displays countdown timers on an LCD, and logs traffic light changes via serial communication. Additionally, it features an interrupt-driven button allowing pedestrians to request crossing, thereby altering the normal traffic sequence.

Key Techniques Used:

Input/Output
Timer
Interrupt
PWM and Servo Control
Serial Communication
LCD Display

2. Functionalities

Traffic Light Control: Manages red, yellow, and green LEDs to simulate traffic signals.
Servo Motor Integration: Adjusts the servo position based on the active traffic light (0° for red, 90° for green).
LCD Display: Shows the current traffic light status and a countdown timer.
Serial Communication: Logs traffic light changes to a serial monitor for monitoring and debugging.
Interrupt-Driven Pedestrian Button: Allows pedestrians to request crossing, triggering a modified traffic sequence.

3. Hardware Components and Schematics

3.1 Components Used

Microcontroller: AVR (e.g., ATmega328P)
LEDs: Red, Yellow, Green
Servo Motor: Standard hobby servo
LCD Display: 16x2 character LCD
Button: Momentary push-button
Resistors: Current-limiting resistors for LEDs
Power Supply: 5V regulated supply
Miscellaneous: Breadboard, jumper wires, USB-to-Serial adapter

3.2 Connection Overview

LEDs:
- Red LED: Connected to PORTB3 through a current-limiting resistor.
- Yellow LED: Connected to PORTB4 through a current-limiting resistor.
- Green LED: Connected to PORTB5 through a current-limiting resistor.
Servo Motor:
- Control Pin: Connected to PORTB2 (OC1B).
- Power: Connected to an external 5V supply with a common ground.
LCD Display:
- Data Lines (D4-D7): Connected to PORTD4-PORTD7.
- Control Lines (RS, EN): Connected to PORTB0 and PORTB1.
Button:
- One Terminal: Connected to PORTD3.
- Other Terminal: Connected to GND with an internal pull-up resistor enabled.
Serial Communication:
- TX Pin (PD1): Connected to RX of USB-to-Serial adapter.
- Common Ground: Ensured between microcontroller and adapter.

4. Techniques and Methodologies

4.1 LED Control

Utilized PORTB pins to control red, yellow, and green LEDs.
Employed bitwise operations to set and clear specific LED states without affecting others.

4.2 Servo Motor Control

Implemented PWM (Pulse Width Modulation) using Timer1 in Fast PWM mode.
Mapped specific angles (0° and 90°) to corresponding pulse widths to position the servo accurately.
Configured the servo to reset to 0° during red light and move to 90° during green light.

4.3 LCD Display

Initialized the LCD in 4-bit mode for efficient data transmission.
Displayed the current traffic light status and a countdown timer on two lines.
Utilized sprintf for formatting countdown strings.

4.4 Serial Communication

Configured USART for serial communication at 9600 baud rate.
Implemented functions to initialize USART, transmit single characters, and send strings.
Logged every traffic light change with appropriate messages for monitoring.

4.5 Interrupt Handling

Configured an external interrupt (INT1) on PORTD3 to detect button presses.
Employed a sequence in the ISR, when the button is pressed, to display the yellow light for 3 seconds and then the green light for 5 seconds for pedestrians crossing the road.

5. Demo Results

5.1 Normal Traffic Sequence

Red Light:
- LED: Red LED on.
- Servo: Positioned at 0°.
- LCD: Displays “RED LIGHT” with an 8-second countdown.
- Serial: Logs “RED LIGHT”.
Green Light:
- LED: Green LED on.
- Servo: Positioned at 90°.
- LCD: Displays “GREEN LIGHT” with a 4-second countdown.
- Serial: Logs “GREEN LIGHT”.
Yellow Light:
- LED: Yellow LED on.
- LCD: Displays “YELLOW LIGHT” with a 2-second countdown.
- Serial: Logs “YELLOW LIGHT”.
  (Repeat cycle)

5.2 Pedestrian Request Sequence

Button Press Detected: Triggers the ISR to execute the special sequence.

Yellow Light:
- LED: Yellow LED on.
- LCD: Displays “Pedestrians” with a 3-second countdown.
- Serial: Logs “YELLOW LIGHT”.
Green Light:
- LED: Green LED on.
- Servo: Positioned at 90°.
- LCD: Displays “GREEN LIGHT” with a 5-second countdown.
- Serial: Logs “GREEN LIGHT”.
Return to Normal Sequence:
- Servo: Resets to 0°.
- Resumes Normal Traffic Cycle.

6. Performance Evaluation

Timing Accuracy: _delay_ms() provided sufficient timing for demonstration purposes, but hardware timers are recommended for higher precision.
Interrupt Responsiveness: The system detected button presses and executed pedestrian requests promptly.
Servo Control: Smooth and accurate transitions between designated angles.
Serial Logging: Reliable logging of traffic light changes for monitoring.
LCD Display: Consistently displayed accurate statuses and countdowns.

Areas for Improvement:

Implement software or hardware debouncing for the button.
Extend servo angle functionality for versatile signaling.
Replace _delay_ms() with timer-based delays for concurrent tasks.
Add error handling for USART transmission and routine interruptions.

7. Conclusion

The Traffic Light Display System integrates various embedded system components controlled by an AVR microcontroller. Interrupt-driven button handling adds interactivity and realism, showcasing practical applications of PWM, serial communication, and user interfaces.

8. References

AVR Microcontroller Datasheet
Arduino Documentation
LCD Display Lab3 Tutorials
Servo Motor Control course material
Serial Communication course material

Posts on Decsnow Media

MIT 6.5940 EfficientML.ai Lab 3: Neural Architecture Search

MIT 6.5940 EfficientML.ai Lab 3: Neural Architecture Search

Introduction

Getting Started: Super Network and the VWW dataset (1 Question, 5 pts)

Question 1 (5 pts): Design space exploration.

Part 1. Predictors (3 Questions, 30 pts)

Question 2 (10 pts): Implement the efficiency predictor.

Question 3 (10 pts): Implement the accuracy predictor.

Question 4 (10 pts): Complete the code for accuracy predictor training.

Part 2. Neural Architecture Search (6 Questions, 65 pts + 10 bonus pts)

Question 5 (5 pts): Complete the following random search agent.

Question 6 (5 pts): Complete the following function.

Question7 (20 pts): Complete the following evolutionary search agent.

Question 8 (10pts): Run evolutionary search and tune evo_params to optimize the results. Describe your findings.

Question 9 (15 pts + 10 pts bonus): Run evolutionary search under real-world constraints.

Question 10 (10 pts): Is it possible to find a subnet with the following efficiency constraints in the current design space?

MIT 6.5940 EfficientML.ai Lab 2: Quantization

MIT 6.5940 EfficientML.ai Lab 2: Quantization

Goals

Contents

Setup

Let’s First Evaluate the Accuracy and Model Size of the FP32 Model

K-Means Quantization

Question 1 (10 pts)

Question 2 (10 pts)

Question 2.1 (5 pts)

Question 2.2 (5 pts)

K-Means Quantization on Whole Model

Trained K-Means Quantization

Question 3 (10 pts)

Linear Quantization

n-bit Integer

Question 4 (10 pts)

Question 5 (15 pts)

Scale

Question 5.1 (3 pts)

zero point

Question 5.2 (4 pts)

Question 5.3 (8 pts)

Special case: linear quantization on weight tensor

Per-channel Linear Quantization

A Quick Peek at Linear Quantization on Weights

Quantized Inference

Question 6 (5 pts)

Quantized Fully-Connected Layer

Question 7 (15 pts)

Quantized Convolution

Question 8 (10 pts)

Question 9 (10 pts)

Question 9.1 (5 pts)

Question 9.2 (5 pts)

Question 10 (5 pts)

MIT 6.5940 EfficientML.ai: Lab 1 Pruning

MIT 6.5940 EfficientML.ai: Lab 1 Pruning

Goals

Contents

Setup

Let’s First Evaluate the Accuracy and Model Size of Dense Model

Let’s see the distribution of weight values

Question 1 (10 pts)

Question 1.1 (5 pts)

Question 1.2 (5 pts)

Fine-grained Pruning

Magnitude-based Pruning

Question 2 (15 pts)

Question 3 (5 pts)

Sensitivity Scan

Question 4 (15 pts)

Question 4.1 (5 pts)

Question 4.2 (5 pts)

Question 4.3 (5 pts)

#Parameters of each layer

Select Sparsity Based on Sensitivity Curves and #Parameters Distribution

Question 5 (10 pts)

Finetune the fine-grained pruned model

Channel Pruning

Remove Channel Weights

Question 6 (10 pts)

Ranking Channels by Importance