How to use PAI-Blade for quantization optimization on GPUs - Platform For AI

PAI-Blade supports INT8 quantization for TensorFlow and PyTorch models on GPU. Quantization reduces model memory usage and increases inference throughput by replacing 32-bit floating-point numbers with lower-bit-width fixed-point integers.

INT8 quantization requires the GPU to support INT8 computation. If the GPU device that you use supports INT8 quantization and quantized models can be accelerated, PAI-Blade performs quantization.

How it works

Online mode vs. offline mode

	Online mode	Offline mode
Calibration dataset	Not required	Required
Acceleration	Lower	Higher (recommended)
Quantization parameters	Computed without a calibration dataset	Computed from calibration data

For TensorFlow models, both modes are supported. For PyTorch models, only offline mode is available — a calibration dataset is required.

Calibration dataset

The calibration dataset should represent the real distribution of your inference inputs. The code examples in this topic use 10 samples for illustration only.

Framework	Format
TensorFlow	A list of `feed_dict` arguments; all values must be `np.ndarray`
PyTorch	A list containing multiple groups of input data

Prerequisites

Before you begin, make sure you have:

A GPU that supports INT8 quantization
PAI-Blade installed in your Python environment
A trained TensorFlow (frozen) or PyTorch (traced) model

Quantize a TensorFlow model

Set optimization_level='o2' in blade.optimize() to enable quantization. For best results, provide a calibration dataset to run in offline mode.

Prepare the calibration dataset:

# Prepare a calibration dataset.
import numpy as np
calib_data = list()
for i in range(10):
    # All values in the feed_dict arguments must be of the np.ndarray type.
    feed_dict = {'input:0': np.ones((32, 224, 224, 3), dtype=np.float32)}
    calib_data.append(feed_dict)

Steps:

Download the sample model, test data, and calibration dataset.

wget https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/test_public_model/bbs/tf_resnet50_v1.5/frozen.pb
wget https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/test_public_model/bbs/tf_resnet50_v1.5/test_bc32.npy
wget https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/test_public_model/bbs/tf_resnet50_v1.5/calib_data_test_bc32.npy

Load the model and data.

import numpy as np
import tensorflow as tf

# Load the model.
graph_def = tf.GraphDef()
with open('frozen.pb', 'rb') as f:
    graph_def.ParseFromString(f.read())

# Load the test data.
test_data = np.load('test_bc32.npy', allow_pickle=True, encoding='bytes',).item()

# Load the calibration dataset.
calib_data = np.load('calib_data_test_bc32.npy', allow_pickle=True, encoding='bytes',)

Run INT8 quantization in offline mode.

import blade
optimized_model, opt_spec, report = blade.optimize(
    model=graph_def,
    optimization_level='o2',
    device_type='gpu',
    test_data=test_data,
    calib_data=calib_data
)

Verify the precision of the quantized model. Use a complete test dataset to check whether precision has dropped significantly. If precision meets your requirements, quantization is complete. If precision loss is too high, enable weight_adjustment to let PAI-Blade automatically adjust model parameters. This option is available for TensorFlow models on GPU only.
```
quant_config = {
    'weight_adjustment': 'true'  # Default value: false.
}
optimized_model, opt_spec, report = blade.optimize(
    model=graph_def,
    optimization_level='o2',
    device_type='gpu',
    test_data=test_data,
    calib_data=calib_data,
    config=blade.Config(quant_config=quant_config)
)
```
For all available quantization configuration options, see blade.Config.

Quantize a PyTorch model

PyTorch models can only be quantized in offline mode. A calibration dataset is required.

Set optimization_level='o2' in blade.optimize() to enable quantization.

Prepare the calibration dataset:

# Prepare a calibration dataset.
import torch
calib_data = list()
for i in range(10):
    image = torch.ones(32, 3, 224, 224)
    calib_data.append(image)

Steps:

Download the sample model, test data, and calibration dataset.

wget https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/test_public_model/bbs/pt_resnet50_v1.5/traced_model.pt
wget https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/test_public_model/bbs/pt_resnet50_v1.5/test_bc32.pth
wget https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/test_public_model/bbs/pt_resnet50_v1.5/calib_data_test_bc32.pth

Load the model and data.

import torch

# Load the model.
pt_model = torch.jit.load('traced_model.pt')

# Load the test data.
test_data = torch.load('test_bc32.pth')

# Load the calibration dataset.
calib_data = torch.load('calib_data_test_bc32.pth')

Run INT8 quantization in offline mode.

import blade
optimized_model, opt_spec, report = blade.optimize(
    model=pt_model,
    optimization_level='o2',
    device_type='gpu',
    test_data=test_data,
    calib_data=calib_data
)

What's next

To optimize a TensorFlow model with other optimization levels, see Optimize a TensorFlow model.
For all quantization configuration options including weight_adjustment, see blade.Config.