xbe.at

Machine learning / Advanced ML topics

Quantization Techniques for Large Machine Learning Models

10 min read, 15 sections

Contents
  1. Introduction to Quantization Techniques
  2. Types of Quantization
  3. Uniform Quantization
  4. Non-Uniform Quantization
  5. Weight Quantization
  6. Activation Quantization
  7. Quantization-Aware Training
  8. Mixed-Precision Quantization
  9. Quantization for Transformer Models
  10. Pruning and Quantization
  11. Quantization for Convolutional Neural Networks
  12. Quantization for Recurrent Neural Networks
  13. Real-life Example: Image Classification
  14. Real-life Example: Natural Language Processing
  15. Additional Resources

Introduction to Quantization Techniques

Quantization is a crucial technique in optimizing large machine learning models, particularly neural networks. It involves reducing the precision of the model's parameters and computations, typically from 32-bit floating-point to lower bit-width representations. This process can significantly reduce model size and improve inference speed with minimal impact on accuracy.

import numpy as np

# Original 32-bit float array
original = np.array([0.1, 0.2, 0.3, 0.4, 0.5], dtype=np.float32)

# Quantize to 8-bit integers
scale = (original.max() - original.min()) / 255
zero_point = -int(original.min() / scale)
quantized = np.round(original / scale + zero_point).astype(np.uint8)

print("Original:", original)
print("Quantized:", quantized)
print("Dequantized:", (quantized.astype(float) - zero_point) * scale)

Types of Quantization

There are several types of quantization, including post-training quantization, quantization-aware training, and dynamic quantization. Post-training quantization is applied after model training, while quantization-aware training incorporates quantization during the training process. Dynamic quantization determines quantization parameters at runtime.

import torch

# Post-training static quantization
model_fp32 = torch.jit.load("model_fp32.pt")
model_int8 = torch.quantization.quantize_dynamic(
    model_fp32,  # the original model
    {torch.nn.Linear},  # a set of layers to dynamically quantize
    dtype=torch.qint8)  # the target dtype for quantized weights

# Save the quantized model
torch.jit.save(model_int8, "model_int8.pt")

Uniform Quantization

Uniform quantization is the simplest form of quantization, where the full range of values is divided into equal-sized intervals. Each value in an interval is mapped to a single quantized value, typically represented by an integer.

def uniform_quantize(x, num_bits=8):
    max_val = np.max(np.abs(x))
    step = 2 * max_val / (2**num_bits - 1)
    return np.round(x / step) * step

# Example usage
x = np.array([-1.2, -0.5, 0, 0.1, 0.9])
x_quantized = uniform_quantize(x)
print("Original:", x)
print("Quantized:", x_quantized)

Non-Uniform Quantization

Non-uniform quantization uses varying interval sizes, often concentrating more intervals in regions with higher frequency or importance. This can lead to better preservation of information in critical areas of the value range.

import numpy as np
import matplotlib.pyplot as plt

def log_quantize(x, num_bits=8):
    sign = np.sign(x)
    x_abs = np.abs(x)
    max_val = np.max(x_abs)
    
    # Log space quantization
    log_x = np.log1p(x_abs / max_val)
    quantized_log = np.round(log_x * (2**num_bits - 1)) / (2**num_bits - 1)
    return sign * max_val * (np.exp(quantized_log) - 1)

# Generate sample data
x = np.linspace(-1, 1, 1000)
y = np.sin(2 * np.pi * x)

# Quantize
y_quantized = log_quantize(y)

# Plot
plt.figure(figsize=(10, 6))
plt.plot(x, y, label='Original')
plt.plot(x, y_quantized, label='Log Quantized')
plt.legend()
plt.title('Log Quantization Example')
plt.show()

Weight Quantization

Weight quantization focuses on reducing the precision of model parameters. This technique can significantly reduce model size and memory bandwidth requirements during inference.

import torch
import torch.nn as nn

class QuantizedLinear(nn.Module):
    def __init__(self, in_features, out_features, bits=8):
        super().__init__()
        self.weight = nn.Parameter(torch.Tensor(out_features, in_features))
        self.bias = nn.Parameter(torch.Tensor(out_features))
        self.bits = bits
        self.scale = nn.Parameter(torch.Tensor(1))
        self.zero_point = nn.Parameter(torch.Tensor(1))
        
        # Initialize parameters
        nn.init.kaiming_uniform_(self.weight)
        nn.init.zeros_(self.bias)
        self.scale.data.fill_(1.0)
        self.zero_point.data.fill_(0.0)

    def forward(self, input):
        # Quantize weights
        qweight = torch.quantize_per_tensor(
            self.weight, self.scale, self.zero_point, torch.qint8
        )
        # Dequantize for computation
        return nn.functional.linear(input, qweight.dequantize(), self.bias)

# Usage
layer = QuantizedLinear(10, 5, bits=8)
x = torch.randn(3, 10)
output = layer(x)
print(output.shape)  # torch.Size([3, 5])

Activation Quantization

Activation quantization reduces the precision of intermediate activations in neural networks. This can lead to reduced memory usage and faster computation during inference.

import torch
import torch.nn as nn

class QuantizedReLU(nn.Module):
    def __init__(self, bits=8):
        super().__init__()
        self.bits = bits

    def forward(self, input):
        # Determine quantization range
        qmin, qmax = 0, 2**self.bits - 1
        scale = (input.max() - input.min()) / (qmax - qmin)
        zero_point = qmin - input.min() / scale

        # Quantize and clip
        output = torch.clamp(torch.round(input / scale + zero_point), qmin, qmax)
        
        # Dequantize
        output = (output - zero_point) * scale

        return torch.clamp(output, min=0)  # ReLU operation

# Usage
activation = QuantizedReLU(bits=4)
x = torch.randn(5, 5)
output = activation(x)
print(output)

Quantization-Aware Training

Quantization-aware training (QAT) involves simulating the effects of quantization during the training process. This allows the model to adapt to quantization, potentially leading to better performance compared to post-training quantization.

import torch
import torch.nn as nn

class QuantizationAwareModule(nn.Module):
    def __init__(self, module, bits=8):
        super().__init__()
        self.module = module
        self.bits = bits

    def forward(self, x):
        if self.training:
            # Simulate quantization noise
            x = self.simulate_quantization(x)
        return self.module(x)

    def simulate_quantization(self, x):
        scale = (x.max() - x.min()) / (2**self.bits - 1)
        x_q = torch.round(x / scale) * scale
        return x_q + (x - x_q).detach()  # STE (Straight-Through Estimator)

# Usage
linear = nn.Linear(10, 5)
qat_linear = QuantizationAwareModule(linear, bits=8)

x = torch.randn(3, 10)
output = qat_linear(x)
print(output.shape)  # torch.Size([3, 5])

Mixed-Precision Quantization

Mixed-precision quantization involves using different bit-widths for different parts of the model. This approach allows for a balance between model size, speed, and accuracy by allocating more bits to sensitive layers and fewer bits to robust layers.

import torch
import torch.nn as nn

class MixedPrecisionModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        self.fc = nn.Linear(32 * 8 * 8, 10)

    def forward(self, x):
        x = torch.quantize_per_tensor(x, scale=1/255, zero_point=0, dtype=torch.quint8)
        x = self.conv1(x.dequantize())
        x = torch.quantize_per_tensor(x, scale=1/127, zero_point=0, dtype=torch.qint8)
        x = self.conv2(x.dequantize())
        x = x.dequantize().view(x.size(0), -1)
        x = self.fc(x)
        return x

# Usage
model = MixedPrecisionModule()
input_tensor = torch.randn(1, 3, 32, 32)
output = model(input_tensor)
print(output.shape)  # torch.Size([1, 10])

Quantization for Transformer Models

Transformer models, which are prevalent in natural language processing tasks, can benefit significantly from quantization. Here's an example of quantizing a transformer encoder layer:

import torch
import torch.nn as nn

class QuantizedTransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward, dropout=0.1, bits=8):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead)
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.bits = bits

    def quantize(self, x):
        scale = (x.max() - x.min()) / (2**self.bits - 1)
        return torch.round(x / scale) * scale

    def forward(self, src, src_mask=None, src_key_padding_mask=None):
        src2 = self.self_attn(src, src, src, attn_mask=src_mask,
                              key_padding_mask=src_key_padding_mask)[0]
        src = src + self.dropout1(src2)
        src = self.norm1(src)
        src2 = self.linear2(self.dropout(torch.relu(self.linear1(src))))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return self.quantize(src)

# Usage
encoder_layer = QuantizedTransformerEncoderLayer(d_model=512, nhead=8, dim_feedforward=2048)
src = torch.rand(10, 32, 512)
out = encoder_layer(src)
print(out.shape)  # torch.Size([10, 32, 512])

Pruning and Quantization

Pruning and quantization are often used together to achieve even greater model compression. Pruning removes less important weights, while quantization reduces the precision of the remaining weights.

import torch
import torch.nn as nn

class PrunedAndQuantizedLinear(nn.Module):
    def __init__(self, in_features, out_features, sparsity=0.5, bits=8):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features)
        self.sparsity = sparsity
        self.bits = bits
        self.mask = None

    def prune(self):
        weight = self.linear.weight.data.abs()
        threshold = weight.view(-1).kthvalue(int(self.sparsity * weight.numel())).values
        self.mask = (weight > threshold).float()

    def quantize(self, x):
        scale = (x.max() - x.min()) / (2**self.bits - 1)
        return torch.round(x / scale) * scale

    def forward(self, x):
        if self.mask is None:
            self.prune()
        weight = self.linear.weight * self.mask
        quantized_weight = self.quantize(weight)
        return nn.functional.linear(x, quantized_weight, self.linear.bias)

# Usage
layer = PrunedAndQuantizedLinear(10, 5, sparsity=0.3, bits=8)
x = torch.randn(3, 10)
output = layer(x)
print(output.shape)  # torch.Size([3, 5])
print(f"Sparsity: {(layer.mask == 0).float().mean().item():.2f}")

Quantization for Convolutional Neural Networks

Convolutional Neural Networks (CNNs) are widely used in computer vision tasks and can benefit greatly from quantization. Here's an example of quantizing a simple CNN:

import torch
import torch.nn as nn
import torch.nn.functional as F

class QuantizedCNN(nn.Module):
    def __init__(self, bits=8):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)
        self.bits = bits

    def quantize(self, x):
        scale = (x.max() - x.min()) / (2**self.bits - 1)
        return torch.round(x / scale) * scale

    def forward(self, x):
        x = self.quantize(self.conv1(x))
        x = F.relu(F.max_pool2d(x, 2))
        x = self.quantize(self.conv2(x))
        x = F.relu(F.max_pool2d(x, 2))
        x = torch.flatten(x, 1)
        x = F.relu(self.quantize(self.fc1(x)))
        x = self.fc2(x)
        return x

# Usage
model = QuantizedCNN()
input_tensor = torch.randn(1, 1, 28, 28)
output = model(input_tensor)
print(output.shape)  # torch.Size([1, 10])

Quantization for Recurrent Neural Networks

Recurrent Neural Networks (RNNs) and their variants like LSTMs and GRUs are crucial for sequence modeling tasks. Quantizing these models can be challenging due to their recurrent nature, but it's still possible and beneficial:

import torch
import torch.nn as nn

class QuantizedLSTM(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, bits=8):
        super().__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.bits = bits

    def quantize(self, x):
        scale = (x.max() - x.min()) / (2**self.bits - 1)
        return torch.round(x / scale) * scale

    def forward(self, x):
        output, (h_n, c_n) = self.lstm(x)
        return self.quantize(output), (self.quantize(h_n), self.quantize(c_n))

# Usage
model = QuantizedLSTM(input_size=10, hidden_size=20, num_layers=2)
input_sequence = torch.randn(32, 5, 10)  # (batch_size, sequence_length, input_size)
output, (h_n, c_n) = model(input_sequence)
print("Output shape:", output.shape)
print("Hidden state shape:", h_n.shape)
print("Cell state shape:", c_n.shape)

Real-life Example: Image Classification

Quantization can significantly reduce the size and improve the inference speed of image classification models. Here's an example using a pre-trained ResNet model:

import torchvision.models as models
import torch.quantization

# Load a pre-trained ResNet model
model = models.resnet18(pretrained=True)

# Set the model to evaluation mode
model.eval()

# Fuse Conv, BN, and ReLU layers
model = torch.quantization.fuse_modules(model, [["conv1", "bn1", "relu"]])

# Specify quantization configuration
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
torch.quantization.prepare(model, inplace=True)

# Calibrate the model (you would normally use a calibration dataset here)
dummy_input = torch.randn(1, 3, 224, 224)
model(dummy_input)

# Convert to quantized model
quantized_model = torch.quantization.convert(model, inplace=True)

# Compare model sizes
print(f"Original model size: {torch.save(model.state_dict(), '/tmp/model.pth')/(1024*1024):.2f} MB")
print(f"Quantized model size: {torch.save(quantized_model.state_dict(), '/tmp/quantized_model.pth')/(1024*1024):.2f} MB")

# Perform inference
input_tensor = torch.randn(1, 3, 224, 224)
output = quantized_model(input_tensor)
print("Output shape:", output.shape)

Real-life Example: Natural Language Processing

Quantization is also beneficial for NLP models, such as those used for sentiment analysis. Here's an example using a simple LSTM-based model:

import torch
import torch.nn as nn

class SentimentAnalysisModel(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim, output_dim):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, output_dim)
        
    def forward(self, text):
        embedded = self.embedding(text)
        output, (hidden, cell) = self.lstm(embedded)
        return self.fc(hidden[-1])

# Create a model
vocab_size = 10000
embed_dim = 100
hidden_dim = 256
output_dim = 2  # binary sentiment (positive/negative)

model = SentimentAnalysisModel(vocab_size, embed_dim, hidden_dim, output_dim)

# Quantize the model
quantized_model = torch.quantization.quantize_dynamic(
    model, {nn.Linear, nn.LSTM, nn.Embedding}, dtype=torch.qint8
)

# Compare model sizes
print(f"Original model size: {torch.save(model.state_dict(), '/tmp/model.pth')/(1024*1024):.2f} MB")
print(f"Quantized model size: {torch.save(quantized_model.state_dict(), '/tmp/quantized_model.pth')/(1024*1024):.2f} MB")

# Perform inference
input_tensor = torch.randint(0, vocab_size, (1, 20))  # (batch_size, sequence_length)
output = quantized_model(input_tensor)
print("Output shape:", output.shape)

Additional Resources

For those interested in diving deeper into quantization techniques for large models, here are some valuable resources:

  1. "Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference" by Jacob et al. (2018) ArXiv: https://arxiv.org/abs/1712.05877
  2. "A Survey of Quantization Methods for Efficient Neural Network Inference" by Gholami et al. (2021) ArXiv: https://arxiv.org/abs/2103.13630
  3. "ZeroQ: A Novel Zero Shot Quantization Framework" by Cai et al. (2020) ArXiv: https://arxiv.org/abs/2001.00281
  4. "HAWQ: Hessian AWare Quantization of Neural Networks with Mixed-Precision" by Dong et al. (2019) ArXiv: https://arxiv.org/abs/1905.03696

These papers provide in-depth discussions on various quantization techniques, their implementations, and their impact on model performance.

More in Advanced ML topics

  1. Role of Cost Functions in Machine Learning11 min

    A cost function, also known as loss function, measures the difference between predicted and actual values in machine learning models. It quantifies the error...

  2. Probability in Machine Learning with Python8 min

    Probability theory forms the backbone of many machine learning algorithms. It provides a framework for dealing with uncertainty and making predictions based...

  3. Second-Order Equations in Machine Learning Algorithms Using Python7 min

    Introduction to Second-Order Equations in Machine Learning and AI

  4. Probability Fundamentals for Machine Learning13 min

    Conditional probability forms the foundation of many machine learning algorithms, particularly in classification tasks. It represents the probability of an...

  5. Self-Supervised Learning with SimCLR in TensorFlow11 min

    SimCLR (Simple Framework for Contrastive Learning of Visual Representations) is a self-supervised learning technique for computer vision tasks. It allows...

  6. Predictive Maintenance and Condition-Based Monitoring Using Machine Learning13 min

    Predictive Maintenance and Condition-Based Monitoring (CBM) are advanced approaches to equipment maintenance that leverage data analytics and machine...