EfficientML.ai Lec 3 - Pruning and Sparsity (Part I) MIT 6.5940, Fall 2023, Zoom
本文是EfficientML.ai Fall 2023课程作业1练习答案，在本次练习里将会对经典的分类神经网络进行剪枝处理，减少模型大小和延迟。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
文章目录
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.
print('Installing torchprofile...')
!pip install torchprofile 1>/dev/null
print('All required packages have been successfully installed!')
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"
random.seed(0)
np.random.seed(0)
torch.manual_seed(0)
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
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
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()
@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()
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
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.
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'])
=> loading checkpoint ‘https://hanlab18.mit.edu/files/course/labs/vgg.cifar.pretrained.pth’
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,
)
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.
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")
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×32
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.
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?
Your Answer:
The weights for the convolutional layers are normally distributed and centered around 0. As we are doing batch normalization, the weights are stable and they are staying close to 0. The input weight convolution layers are more dispersed around 0 because the input layer are pixel values ranging from 0-255, so the weights are adapting for unnormalized weights (some extreme peaks). But for all the intermediate layers, we see a peak at 0 and gradual distribution on either side. For the last layer, the weights are negatively skewed.
Question 1.2 (5 pts)
How do these characteristics help pruning?
Your Answer:
If we prune our model around the 0 weight value, we will prune many values as the weights are following standard normal distribution. As many weights are nearly 0, hence we can prune them as they don’t affect the final output very much. The extreme weights affect the outut a lot, hence we can prune the approximately zero weights