Optimizing Deep Learning Models with Pruning: A Practical Guide

Exploring and Implementing Pruning Methods with TensorFlow and PyTorch

Optimizing Deep Learning Model with Pruning: A Practical Guide

If you’re interested in improving the efficiency and complexity of your machine and deep learning models, you may have heard of a technique called model pruning.

In this blog post, I’ll explore what model pruning is, the different types of pruning available, and how to implement it using TensorFlow and PyTorch.

Model pruning is a powerful tool for optimizing AI models, and it can be especially useful for deploying models on resource-constrained devices.

So if you want to learn more about this technique and how to use it in your own projects, keep reading!

Definition and Background
Pruning Techniques
Implementing Model Pruning with TensorFlow
Implementing Model Pruning with PyTorch
Key Takeaways, Final Notes, and References

Definition and Background

Definition of Model Pruning

In general, model pruning refers to the technique for improving the efficiency and complexity of machine learning models by removing unnecessary parameters.

It is a strong tool for the optimization of large machine and deep learning models, and can be particularly valuable for deploying models on resource-constrained devices.

In this section, I’ll explore the different types of model pruning, the motivations for using it, and some of the challenges and limitations of the technique.

Motivation for Model Pruning

There are several motivations for using model pruning, including:

Reducing the size and complexity of machine learning models: By removing unnecessary parameters, model pruning can assist in reducing the size and complexity of a model, making it easier to deploy and run on resource-constrained devices.
Improving performance: Model pruning can help to improve the performance of a machine learning model by reducing overfitting and improving generalization. It can also lead to faster training times, as fewer parameters need to be updated during training.
Increasing interpretability: By removing unnecessary parameters, it might also increase the interpretability of a machine learning model, making its decision-making process easier to understand.

Types of Model Pruning

There are two main types of model pruning: weight pruning and neuron pruning.

Weight Pruning:

The idea behind weight pruning involves removing individual weights or connections within a neural network that are not contributing significantly to the model’s performance. [1]

This can be done through techniques such as magnitude-based pruning, where weights with small magnitudes are removed, or gradient-based pruning, where weights with small gradients are removed.

Applying weight pruning can be an effective way to reduce the complexity of a model, but it can also lead to a loss of accuracy if too many weights are removed.

Neuron Pruning:

Neuron pruning, also known as structure pruning, involves removing entire neurons or layers from a neural network. [2]

That can be performed by techniques such as low-density pruning, in which neurons with low activity are removed, or filter pruning, where filters with low importance are removed.

Similarly to weight pruning neuron pruning can be an effective way to reduce the size of a model, but can also lead to a loss of accuracy if too many neurons or layers are removed.

Structured vs Unstructured Pruning:

Structured pruning and unstructured pruning are two different approaches to model pruning that can be used to remove unnecessary parameters. [3]

Structured pruning involves removing parameters in a way that preserves the structure of the model, for instance by removing weights or neurons in a regular pattern.

This can be contrasted with unstructured pruning, which involves removing parameters in an irregular pattern, like removing weights or neurons randomly or based on some measure of importance.

Challenges and Limitations of Model Pruning

Although model pruning can be a helpful and strong tool for optimizing machine learning models, there are also some challenges and limitations to consider. These include:

Balancing complexity and performance: It is important to carefully balance the complexity of a model with its performance, as removing too many parameters can lead to a loss of accuracy.
Over-pruning: If a model is over-pruned, it can lead to a loss of accuracy and poor performance. This can be especially challenging for deep neural networks, where it can be difficult to determine which parameters are necessary.
Lack of interpretability: In some cases, model pruning can lead to a lack of interpretability, as the relationships between the remaining parameters may be more complex and difficult to understand.

Overall, model pruning is a useful technique for improving the efficiency and complexity of machine learning models, but it is important to carefully consider the trade-offs and limitations of the technique to achieve the best results.

Pruning Techniques

As outlined in the previous section, several techniques can be used for model pruning, including weight pruning and neuron pruning.

In this section, I will go into more detail describing the different approaches to these techniques and discuss some of the other pruning methods that have been proposed in the literature.

Weight Pruning

Weight pruning involves removing individual weights or connections within a neural network that are not contributing significantly to the model’s performance.

For instance, magnitude-based pruning is a method, where weights with small magnitudes are removed. Another technique is gradient-based pruning, where weights with small gradients are discarded.

Magnitude-based pruning is a simple and intuitive approach to weight pruning, as it removes weights based on their magnitude relative to the other weights in the model. It is typically implemented by setting a threshold for the minimum magnitude of a weight, and any weights below this threshold are discarded.

However, one has to keep in mind that magnitude-based pruning can be sensitive to the choice of threshold, and it may remove important weights if the threshold is set too high.

On the other hand, gradient-based pruning is another approach to weight pruning that erases weights based on their gradient during training. It is classically implemented by setting a threshold for the minimum gradient of a weight, and any weights below this threshold are removed.

While gradient-based pruning can be more robust than magnitude-based pruning, as it takes into account the influence of the weight on the model’s performance it can be computationally intensive, as it requires the computation of gradients during training.

All in all, weight pruning might lead to a loss of accuracy, however, if conducted carefully it is an effective way to reduce the complexity of a model.

Neuron Pruning

Neuron or structured pruning involves removing entire neurons or layers from a neural network. As already outlined above, this can be done through methods like low-density pruning, where neurons with low activity are removed, or filter pruning, where filters with low importance are erased.

Low-density pruning is a simple and intuitive approach to neuron pruning, as it removes neurons based on their activity relative to the other neurons in the model.

Next to that filter pruning is another approach to neuron pruning that removes filters with low importance from a convolutional neural network (CNN). Typically it is implemented by ranking the filters based on some measure of importance, such as the magnitude of their weights or their contribution to the model’s performance, and removing the lowest ranking filters.

In comparison to low-density pruning, it can be more targeted, as it focuses on specific filters rather than entire neurons. Having said that, it can be more computationally intensive, as it requires the computation of importance scores for each filter.

Summed up, neuron pruning can be an effective way to reduce the size of a model, if done right. Else it can also lead to a loss of accuracy if too many neurons or layers are removed.

Comparison of Weight Pruning and Neuron Pruning

Weight pruning and neuron pruning are two different approaches to model pruning that can be used to reduce the complexity and size of a machine learning model, respectively.

Weight pruning involves removing individual weights or connections within a neural network, whereas neuron pruning involves removing entire neurons or layers.

In general, weight pruning is more targeted than neuron pruning, as it allows for the removal of specific weights rather than entire neurons. However, it can also be more sensitive to the choice of pruning threshold and may result in a loss of accuracy if too many weights are taken away.

Neuron pruning, on the other hand, can be more robust, as it erases entire neurons rather than individual weights. However, it can also be more aggressive, as it removes entire neurons or layers rather than just individual weights, and may result in a greater loss of accuracy if too many neurons are removed.

When choosing between weight pruning and neuron pruning, it is important to consider the specific goals and constraints of your machine learning project.

If you are primarily interested in reducing the complexity of your model, weight pruning may be a good choice. But if you are mainly focused on reducing the size of your model, neuron pruning may be a better choice.

It is also important to carefully balance the complexity and size of your model with its performance, as removing too many parameters can lead to a loss of accuracy.

Structured vs Unstructured Model Pruning

Comparison between structured and unstructured pruning

In addition to the different types of model pruning, such as weight pruning and neuron pruning, there are also two main approaches to pruning [4]:

Structured Pruning
Unstructured Pruning.

Structured pruning involves removing parameters in a way that preserves the structure of the model, such as by removing weights or neurons in a regular pattern.

Unstructured pruning, on the other hand, involves removing parameters in an irregular pattern, such as by removing weights or neurons randomly or based on some measure of importance.

Structured pruning is a more efficient and easier-to-implement method than unstructured pruning, but it may be less effective at reducing the complexity of the model.

If one wants to have a more effective at reducing the complexity of the model unstructured pruning is the better choice, but it can be more computationally intensive and difficult to implement.

As with the pruning type when choosing between structured and unstructured pruning, the specific goals and constraints of your project have to be considered.

If efficiency and ease of implementation are important, structured pruning may be a good choice. If maximum complexity reduction is the primary goal, unstructured pruning may be a better pick.

Other Pruning Techniques

In addition to weight pruning and neuron pruning, several other pruning techniques have been proposed in the literature. These include:

Feature map pruning: Feature map pruning [5] involves removing feature maps from a CNN, rather than individual weights or neurons. It can be implemented using techniques similar to those used for weight or neuron pruning, such as magnitude-based or low-density pruning.
Channel pruning: Channel pruning [6, 7] is a variant of feature map pruning that removes channels rather than individual feature maps. It can be implemented using techniques similar to those used for weight or neuron pruning, such as magnitude-based or low-density pruning.
Iterative pruning: Iterative pruning [8] is a technique that involves repeatedly pruning a model and fine-tuning it until a desired level of complexity is achieved. It can be implemented using any of the pruning techniques discussed above, and it allows for a more gradual and controlled reduction in model complexity.

Implementing Model Pruning with TensorFlow

Here I will show a simple example of model pruning using the TensorFlow framework [9]. The full notebook can be found here.

Keep in mind that my intention is not to define the best possible deep neural net structure with an optimized training pipeline.

First, import the necessary libraries:

# Import necessary libraries
import tensorflow as tf
from tensorflow import keras
import tensorflow_model_optimization.sparsity.keras as sp_keras

Continue loading and preprocessing the dataset, here I will simply choose the MNIST dataset:

# Load the dataset
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

# Normalize the input data
x_train = x_train / 255.0
x_test = x_test / 255.0

Next, define and compile a simple convolutional neural network:

# Define the neural network
model = keras.Sequential([
  keras.layers.InputLayer(input_shape=(28, 28)),
  keras.layers.Reshape(target_shape=(28, 28, 1)),
  keras.layers.Conv2D(filters=12, kernel_size=(3, 3), activation='relu'),
  keras.layers.MaxPooling2D(pool_size=(2, 2)),
  keras.layers.Flatten(),
  keras.layers.Dense(10, activation='softmax')
])

# Compile the model
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

After the model is compiled, fit and evaluate it.

# Train the model
model.fit(x_train, y_train, epochs=20)

# Evaluate the model
test_loss, test_acc = model.evaluate(x_test, y_test)
print('Test accuracy:', test_acc)

Next, use the `prune_low_magnitude()` function to prune the model using a polynomial decay schedule. This function removes weights with low magnitude, which can be effective for reducing the complexity of the model.

# Prune the model
pruning_params = {
  'pruning_schedule': sp_keras.PolynomialDecay(
      initial_sparsity=0.50, final_sparsity=0.80,
      begin_step=2000, end_step=4000)
}

model_for_pruning = sp_keras.prune_low_magnitude(model, **pruning_params)

# Compile the pruned model
model_for_pruning.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

Train and evaluate the pruned model:

# Train the pruned model
model_for_pruning.fit(x_train, y_train, epochs=20,
                      callbacks=[UpdatePruningStep()])

# Evaluate the pruned model
test_loss, test_acc = model_for_pruning.evaluate(x_test, y_test)
print('Test accuracy:', test_acc)

So in my case, the test accuracy decreased from 98.2 % to 97.0 % while decreasing the model size by 80 % in total.

Keep in mind that this is not a fully optimized example.

Implementing Model Pruning with PyTorch

Now I will show a similar example of model pruning using the PyTorch framework [10]. The complete notebook is accessible here.

Again, my intention is not to define the best possible deep neural net structure with an optimized training pipeline.

First, import the necessary libraries:

# Import necessary libraries
import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torch.nn.utils.prune as prune

from torchvision import transforms
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader

Continue loading and preprocessing the dataset, here I will go for the MNIST dataset again:

# Define the transform for the data
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# Load the dataset
trainset = MNIST(root='./data', train=True, download=True, transform=transform)
testset = MNIST(root='./data', train=False, download=True, transform=transform)

# Define the dataloaders
trainloader = DataLoader(trainset, batch_size=64, shuffle=True)
testloader = DataLoader(testset, batch_size=64, shuffle=True)

Next, define a simple convolutional neural network:

# Define the CNN model
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.fc1 = nn.Linear(7 * 7 * 64, 1024)
        self.fc2 = nn.Linear(1024, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2, 2)
        x = x.view(-1, 7 * 7 * 64)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = CNN()

# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Define the optimizer and loss function
optimizer = optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()

Run the training pipeline:

# Train the model
for epoch in range(20):
    running_loss = 0.0
    for i, data in enumerate(trainloader):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        
        loss = criterion(outputs, labels)
        loss.backward()
        
        optimizer.step()
        running_loss += loss.item()
        
        if i % len(trainloader) == len(trainloader) - 1:
            print(f'Epoch {epoch + 1}: loss: {round(running_loss / len(trainloader), 3)}')
            running_loss = 0.0

Now evaluate the model (code will be skipped here). After that, you can define the model pruner as follows, here global unstructured pruning is applied to the CNN:

parameters_to_prune = (
    (model.conv1, 'weight'),
    (model.conv2, 'weight'),
    (model.fc1, 'weight'),
    (model.fc2, 'weight')
)

prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.9,
)

Evaluating the model shows a performance decrease from 99.0 % down to 98.8 %. Similar to the TensorFlow example, you can now fine-tune the pruned model, which is keeping a sparsity level of 90 %.

Doing that yields an even better performance than the unpruned model with about 99.26 % test performance.

Key Takeaways, Final Notes, and References

Key Takeaways

Summing the findings of this blog up provides the following insights:

It is possible to significantly reduce the complexity of a trained model through pruning while maintaining or even improving its performance.
It is important to carefully select the pruning method and hyperparameters to ensure that the pruned model maintains its performance. This may involve using iterative or weight pruning, setting a target sparsity level, or using techniques such as weight re-initialization or fine-tuning to optimize the pruned model.
It is crucial to consider the trade-off between model complexity and performance when deciding whether to prune a model. In some cases, the benefits of pruning may not outweigh the potential loss of performance, and it may be more advantageous to keep a larger but more accurate model.
It is of utmost importance to take into account the specific use case and deployment environment when deciding whether to prune a model. In resource-constrained environments, such as mobile devices or edge devices, the reduction in complexity and improvement in efficiency offered by pruning may be particularly valuable.
Modern deep learning frameworks provide easy access to model pruning

Final Notes

In summary, model pruning is a powerful technique for reducing the complexity of machine learning models while maintaining or even improving their performance.

By carefully selecting the pruning method and hyperparameters, and considering the trade-off between complexity and performance, it is possible to significantly reduce the size and computational requirements of a trained model without sacrificing its accuracy.

Model pruning can be applied to a wide range of tasks and models, and is particularly valuable in resource-constrained environments where the speed and size of the model are important considerations.