目录

Plugin-based Quantization

Preparing the Dependency Environment

Plugin-based quantization depends on the environment of your training project. Currently, the size reduction tool supports two frameworks for plugin-based quantization: TensorFlow and PyTorch.

After the environment is configured for plugin-based quantization, quantization can be performed in either non-training or retraining mode.

TensorFlow Plugin-based Quantization

Perform the following steps for tuning a TensorFlow model:

Prerequisites: The model training datasets and full-precision base model (.ckpt file) have been prepared.

  1. Prepare the TensorFlow environment (Preparing the TensorFlow Environment).

  2. Call the API to generate a strategy template for model tuning (Obtaining the Template for Quantization Configuration).

  3. Configure the tuning strategy file (Configuring the Tuning Strategy).

  4. Train the model (Training or Calibrating the Model). (Models are trained in retraining mode and calibrated in non-training mode.)

  5. Extract the model and quantization parameters (Converting the Model).

[h2]Preparing the TensorFlow Environment

Prepare the Linux environment. To use the plugin-based quantization function of the size reduction tool, you need to prepare the following dependencies.

[h2]Obtaining the Template for Quantization Configuration

Add dopt_tf_py3 in the decompressed DDK to the Python environment variable:

import sys
sys.path.append(".../dopt_tf_py3") ## The path should be an absolute path.

Call the API to generate a JSON quantization configuration file for the TensorFlow model to be quantized.

from dopt.dopt_tf.opt_main import generate_config_file
generate_config_file(sess, dst_path="./config_gen.json")
Parameter Mandatory Description
sess Yes sess is used to load the original floating-point TensorFlow model.
dst_path Yes dst_path indicates the path of the generated JSON configuration file.

The generated configuration is as follows:

{
    "quant_strategy_for op 'Conv2D'": [
        "Quant_INT8-8"
    ],
    "quant_strategy_for op 'MatMul'": [
        "Quant_INT8-8"
    ],
    "layer_strategy": {
        "model/resnet_model/conv2d/Conv2D": {
            "op_type": "Conv2D",
            "quant_strategy": "float"
        },
        "model/resnet_model/conv2d_1/Conv2D": {
            "op_type": "Conv2D",
            "quant_strategy": "float"
        },

    }
}

[h2]Configuring the Tuning Strategy

Determine the layer to be quantized based on the configuration file generated above, and change the value of quant_strategy from float to the quantization strategy (for example, Quant_INT8-8).

{
    "quant_strategy_for op 'Conv2D'": [
        "Quant_INT8-8"
    ],
    "quant_strategy_for op 'MatMul'": [
        "Quant_INT8-8"
    ],
    "layer_strategy": {
        "model/resnet_model/conv2d/Conv2D": {
            "op_type": "Conv2D",
            "quant_strategy": "Quant_INT8-8"
        },
        "model/resnet_model/conv2d_1/Conv2D": {
            "op_type": "Conv2D",
            "quant_strategy": "Quant_INT8-8"
        },

    }
}

  • layer_strategy contains all layers that can be quantized in the current version. Please do not add any layers additionally.

  • You need to select a supported quantization strategy based on op_type. If op_type is set to an unsupported quantization strategy, errors will be reported in subsequent quantization phases.

[h2]Training the Model

To perform plugin-based quantization, add the path of the dopt_tf_py3 file in the decompressed DDK to the Python environment variable:

import sys
sys.path.append(".../dopt_tf_py3") ## The path should be an absolute path.

from dopt.dopt_tf.opt_main import optimize_model

Call the optimize_model API to construct a model in session and quantize the model based on the configured JSON quantization strategy.

Train the model, and determine whether the model has met the expected results by checking the loss of the loss function. If so, stop the training and save the checkpoint (CKPT) data.

## Save the corresponding CKPT data.
saver = tf.train.Saver()
saver.save(sess, train_ckpt_path)

To perform non-training quantization, call the set_calibrate_state API to enable the non-training mode.

from dopt.dopt_tf.opt_main import set_calibrate_state

calibration_mode = True
set_calibrate_state(sess, calibration_mode )
Parameter Mandatory Description
sess Yes Quantized TensorFlow session.
calibration_mode Yes

Boolean variable, indicating whether to enable the non-training (calibration) mode.

- True: Enable the non-training (calibration) mode.

- False: Disable the non-training (calibration) mode.

After setting the non-training mode, use the quantized session (sess) to perform forward inference on the calibration set. Then you will be able to start non-training quantization.

[h2]Converting the Model

Collect quantization parameters after you have completed model quantization and retraining, or completed forward inference on the calibration set. Extract the quantization parameters from the session to the quantization configuration file, and generate the corresponding PB file.

Add the path of the dopt_tf_py3 file in the decompressed DDK to the Python environment variable.

import sys
sys.path.append(".../dopt_tf_py3") ## The path should be an absolute path.

from dopt.dopt_tf.opt_main import generate_final_model

Initialize the original floating-point model, load the CKPT file used for model quantization and training, and use the provided API to extract parameters.

with tf.Session(config=config) as sess:
    build_tf_model()
        generate_final_model(
        sess,
        config_file         = FLAGS.config_file,
        output_name_list    = output_name_list,
        ckpt_file           = train_ckpt_path,
        output_dir          = output_dir
    )

build_tf_model: model graph to be quantized in session. You can define it as needed.

Parameter Mandatory Description
sess Yes sess is used to load the original floating-point TensorFlow model.
config_file Yes Path of the JSON file for the layer to be quantized, which should be the same as the settings when you call optimize_model. You need to set this parameter manually.
output_name_list Yes Name of the output node at the last layer of the PB model. The data type is list.
ckpt_file Yes CKPT data saved during retraining or calibration.
output_dir No Path for saving the extracted PB file and quantization parameters. Make sure that you set a valid path. If this parameter is not set, the model will be saved in the current path.

PyTorch Plugin-based Quantization

Perform the following steps for tuning a PyTorch model:

Prerequisites: The model training datasets and full-precision base model (.pth file) have been prepared.

  1. Prepare the PyTorch environment (Preparing the PyTorch Environment).

  2. Call the API to generate a JSON quantization configuration file (Generating a JSON Quantization Configuration File).

  3. Modify the strategy file configuration (Configuring the Tuning Strategy).

  4. Call the API to train the quantized model (Training or Calibrating the Model).

  5. Call the API to generate the quantized ONNX model and quantization parameter file (Generating a Quantized Model).

[h2]Preparing the PyTorch Environment

To use the plugin-based quantization function of the size reduction tool, you need to prepare the following dependencies.

  • The path can contain uppercase letters, lowercase letters, digits, and underscores (_).

  • The file name can contain uppercase letters, lowercase letters, digits, underscores (_), and periods (.).

[h2]Generating a JSON Quantization Configuration File

Add the path of the dopt_pytorch_py3 file in the decompressed DDK to the Python environment variable.

import sys
sys.path.append(".../dopt_pytorch_py3") ## The path should be an absolute path.

from dopt.dopt_torch.opt_main import generate_config_file

Call the API to generate a JSON quantization configuration file for the PyTorch model to be quantized.

generate_config_file(model, input_shape, dst_path="./config_gen.json") # model: torch.nn.Module, input_shape : "input1:input1.shape;input2:input2.shape"
Parameter Mandatory Description
model Yes This parameter is of the torch.nn.Module type, indicating models not encapsulated by distributed APIs such as torch.nn.parallel.DistributedDataParallel.
input_shape Yes Shape of the final model to be deployed, not the shape of the training data in the training phase.
dst_path Yes Path of the generated JSON configuration file.

[h2]Configuring the Tuning Strategy

The generated configuration is as follows:

{
    "quant_strategy_for op '<class 'torch.nn.modules.conv.Conv2d'>'": [
        "Quant_INT8-8"
    ],
    "quant_strategy_for op '<class 'torch.nn.modules.linear.Linear'>'": [
        "Quant_INT8-8"
    ],
    "input_shape": "x:1,3,224,224",
    "params_version": "v1",
    "layer_strategy": {
        "conv1/Conv2D": {
            "op_type": "<class 'torch.nn.modules.conv.Conv2d'>",
            "quant_strategy": "float"
        },
        "conv2/Conv2D": {
            "op_type": "<class 'torch.nn.modules.conv.Conv2d'>",
            "quant_strategy": "float"
        },

    }
}

You can modify the quantization configuration as required.

Taking the preceding JSON configuration as an example, you can modify it as follows:

{
    "quant_strategy_for op '<class 'torch.nn.modules.conv.Conv2d'>'": [
        "Quant_INT8-8"
    ],
    "quant_strategy_for op '<class 'torch.nn.modules.linear.Linear'>'": [
        "Quant_INT8-8"
    ],
    "input_shape": "x:1,3,224,224",
    "params_version": "v1",
    "layer_strategy": {
        "conv1/Conv2D": {
            "op_type": "<class 'torch.nn.modules.conv.Conv2d'>",
            "quant_strategy": "Quant_INT8-8"
        },
        "conv2/Conv2D": {
            "op_type": "<class 'torch.nn.modules.conv.Conv2d'>",
            "quant_strategy": "float"
        },

    }
}

[h2]Training the Model

Add the path of the dopt_pytorch_py3 file in the decompressed DDK to the Python environment variable.

import sys
sys.path.append(".../dopt_pytorch_py3") ## The path should be an absolute path.
from dopt.dopt_torch.opt_main import optimize_model

You need to call APIs for model quantization in the phase of defining the PyTorch model.

quant_model = optimize_model(model, config_path) ## **config_path** indicates the JSON file's absolute path you have modified in the previous step.
Parameter Mandatory Description
model Yes This parameter is of the torch.nn.Module type, indicating models not encapsulated by distributed APIs such as torch.nn.parallel.DistributedDataParallel.
config_path Yes Shape of the final model to be deployed, not the shape of the training data in the training phase.

You can call APIs as required to obtain quantization-related loss functions and add them to the tuning and training job (optional and not recommended for simple jobs).

from dopt.dopt_torch.opt_main import get_quant_loss
quant_loss = get_quant_loss(quant_model)
total_loss = loss + quant_weight * quant_loss ## **quant_weight** is a hyperparameter of the quantization loss.
Parameter Mandatory Description
quant_model Yes Quantized model after the optimize_model API is called.

Save the checkpoint of the quantized and trained model.

torch.save(quant_model.state_dict(),"quant.pth")

If retraining is not needed, you can perform quantization only in non-training mode. To enable non-training mode, call the set_calibrate_state API.

from dopt.dopt_torch.opt_main import set_calibrate_state
calibrate_mode = True
set_calibrate_state(quant_model, calibrate_mode)
Parameter Mandatory Description
quant_model Yes Quantized PyTorch model.
calibration_mode Yes

Boolean variable, indicating whether to enable the non-training (calibration) mode.

- True: Enable the non-training (calibration) mode.

- False: Disable the non-training (calibration) mode.

After setting the non-training mode, use the quantized model (quant_model) to perform forward inference on the calibration set. Then you will be able to start non-training quantization.

  • When calling optimize_model, make sure that the input model is not encapsulated by distributed APIs such as torch.nn.parallel.DistributedDataParallel, and that the optimizer is used to optimize the quant_model instead of the original floating-point model.

  • The quantization loss is optional. You need to adjust quant_loss based on the actual quantization loss and the original loss, by virtue of hyperparameter tuning. Generally, the proportion between the quantization loss function and the original loss function is 1:20. You can reduce the proportion of the quantization loss, in case of great impact on the optimization.

[h2]Generating a Quantized Model

Add the path of the dopt_pytorch_py3 file in the decompressed DDK to the Python environment variable.

import sys
sys.path.append(".../dopt_pytorch_py3") ## The path should be an absolute path.

from dopt.dopt_torch.opt_main import generate_final_model

After the training is complete, use the corresponding API to generate the final model.

generate_final_model(
    model,              ## Floating-point model.
    config_file,         ## JSON file for quantization configuration.
    pth_file = "quant.pth",    ## Quantized .pth file.
    output_dir = "./"   ## Folder for storing the generated quantization parameters and the PyTorch model.
)
Parameter Mandatory Description
model Yes Floating-point model.
config_file Yes Path of the JSON configuration file, which should be the same as config_path set when you call the optimize_model API.
pth_file Yes Weight file for quantization and training.
output_dir Yes Path for storing the PyTorch model and quantization parameters. Make sure that you set a valid path. If this parameter is not set, the model will be saved in the current path by default.

  • The input model for calling the generate_final_model function must be a rebuilt floating-point model.

  • The generated PyTorch model does not contain quantization parameters, but some graph fusion and fake-quantization processing are performed.