六. 部署分类器-int8-calibration

前言

自动驾驶之心推出的 《CUDA与TensorRT部署实战课程》，链接。记录下个人学习笔记，仅供自己参考

本次课程我们来学习课程第六章—部署分类器，一起来学习分类模型的 INT8 量化

课程大纲可以看下面的思维导图

0. 简述

本小节目标：学习如何用 C++ 实现 INT8 的 calibrator

这个小节我们主要讲 PTQ 量化，看看 calibration 是怎么做的

下面我们开始本次课程的学习🤗

1. 案例运行

在正式开始课程之前，博主先带大家跑通 6.3-int8-calibration 这个小节的案例🤗

源代码获取地址：https://github.com/kalfazed/tensorrt_starter

首先大家需要把 tensorrt_starter 这个项目给 clone 下来，指令如下：

git clone https://github.com/kalfazed/tensorrt_starter.git

也可手动点击下载，点击右上角的 Code 按键，将代码下载下来。至此整个项目就已经准备好了。也可以点击 here 下载博主准备好的源代码（注意代码下载于 2024/7/14 日，若有改动请参考最新）

由于 INT8 量化需要 calibration dataset 校准数据集，所以我们还需要下载 ImageNet 数据集，一般来说校准图片选个 1000 张左右就足够了，因此我们没有必要把整个 ImageNet 数据集下载下来，这个数据量太庞大了，我们只要在它的训练集中挑选个几百上千张图片做校准就足够了

我们可以在 kaggle 上面下载 ImageNet 数据集，大家可以点击 here 查看：

这里博主选择的是 part0 部分，点击进去：

我们可以看到右上角有个 Download(78G) 选项，点击它就利用下载整个 part0 部分的数据集，但是这个数据太大了不是我们想要的，我们看能不能只下载部分数据，我们往下找会发现：

从上图中我们可以知道整个数据由三部分组成

• idx_files：索引文件，用来映射数据集中的数据点到它们在对应数据文件中的位置
• train：训练图片文件，采用 TensorFlow Record（TFRecord）格式存储，并不是单独的 .jpg 或 .png 图片
• validataion：验证图片文件，同样采用 TFRecord 格式存储

我们点击 train 会发现它其实由很多个部分组成，我们下载其中的部分数据即可，这里博主直接选择 train-00000-of-01024 进行下载：

下载完后是一个 .zip 的压缩包，我们解压得到 train-00000-of-01024 数据文件，接着我们需要将它恢复成原始的图片，实现代码如下：(from ChatGPT)

import os
import tensorflow as tf

# 设置你的 TFRecord 文件路径和输出目录
tfrecord_path = 'train-00000-of-01024'
output_folder = 'images'
labels_output_file = 'labels.txt'

# 确保输出目录存在
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

# 解析函数，用于从 TFRecord 中解析图像数据和标签
def _parse_function(proto):
    # 定义解析的字典
    keys_to_features = {
        'image/encoded': tf.io.FixedLenFeature([], tf.string),
        'image/class/label': tf.io.FixedLenFeature([], tf.int64)  # 假设标签字段是这样的
    }
    # 解析一条记录
    parsed_features = tf.io.parse_single_example(proto, keys_to_features)
    # 解码 JPEG 图像
    image = tf.image.decode_jpeg(parsed_features['image/encoded'])
    label = parsed_features['image/class/label']
    return image, label

# 读取 TFRecord 文件
dataset = tf.data.TFRecordDataset(tfrecord_path)
dataset = dataset.map(_parse_function)  # 解析数据

# 准备保存图像和标签
with open(labels_output_file, 'w') as f:
    for i, (image, label) in enumerate(dataset):
        image_path = os.path.join(output_folder, f'image_{i}.JPEG')
        # 保存图像
        image = tf.io.encode_jpeg(image)
        tf.io.write_file(image_path, image)
        # 将标签写入文件
        f.write(f'{image_path}: {label.numpy()}\n')
        print(f'Saved {image_path}, label: {label.numpy()}')

执行后输出如下：

在当前目录下会生成 images 文件夹，里面是恢复出来的图片总共 1251 张，同时还有一个 label.txt 文件里面保存着每张图片对应的 label 标签，如下图所示：

大家可以自行下载，也可以点击 here 下载博主准备好的校准数据集

校准数据集准备好后我们还要准备一个 .txt 文档，该文档里面保存着每张校准图片的完整路径，因为本小节案例代码中是通过解析 calibration/calibration_list_imagenet.txt 去读取每张校准图片进行量化的，因此我们需要把我们自己校准数据集的路径覆盖写入到 calibration_list_imagenet.txt 中

在 6.3-int-calibration 主目录下新建一个 python 脚本文件，内容如下：(from ChatGPT)

import os
import re

# 设置图像文件夹的路径
image_folder = '/home/jarvis/Learn/tensorrt_starter/chapter6-deploy-classification-and-inference-design/6.3-in8-calibration/images'
# 设置输出文件的路径
output_file = 'calibration/calibration_list_imagenet.txt'

# 辅助函数，用于从文件名中提取数字
def extract_number(filename):
    # 从文件名中提取数字，确保正则表达式匹配 "任何数字" 部分
    match = re.search(r'\d+', filename)
    # 如果找到数字，转换为整数；否则返回一个很大的数字，以便将这些文件放在列表末尾
    return int(match.group()) if match else float('inf')

# 初始化一个列表用来存储图像路径
image_paths = []

# 遍历指定文件夹
for filename in os.listdir(image_folder):
    # 检查文件是否为 JPEG 图像
    if filename.endswith('.JPEG') or filename.endswith('.jpg'):
        # 构建完整的文件路径
        file_path = os.path.join(image_folder, filename)
        # 添加到列表
        image_paths.append(file_path)

# 对图像路径列表按文件名中的数字排序
image_paths.sort(key=lambda path: extract_number(os.path.basename(path)))

# 打开文件准备写入
with open(output_file, 'w') as f:
    for path in image_paths:
        # 将路径写入文件
        f.write(path + '\n')

print(f'Image paths have been saved to {output_file}, number = {len(image_paths)}')

注意将 image_folder 的路径修改为你自己的校准数据集路径，执行上述脚本输出如下：

我们在 calibration/calibration_list_imagenet.txt 中可以看到路径的修改，如下图所示：

至此，我们完成了校准图片的准备工作，下面我们就来执行这个小节的案例

整个项目后续需要使用的软件主要有 CUDA、cuDNN、TensorRT、OpenCV，大家可以参考 Ubuntu20.04软件安装大全 进行相应软件的安装，博主这里不再赘述

假设你的项目、环境准备完成，下面我们一起来运行下 6.3-int8-calibration 小节案例代码

开始之前我们需要创建几个文件夹，在 tensorrt_starter/chapter6-deploy-classification-and-inference-design/6.3-int8-calibration 小节中创建一个 models 文件夹，接着在 models 文件夹下创建一个 onnx 和 engine 文件夹，总共三个文件夹需要创建

创建完后 6.3 小节整个目录结构如下：

接着我们需要执行 python 脚本创建一个 ONNX 模型，先进入到 6.3 小节中：

cd tensorrt_starter/chapter6-deploy-classification-and-inference-design/6.2-deploy-classification-advanced

执行如下指令：

python src/python/export_pretained.py -d ./models/onnx/

Note：大家需要准备一个虚拟环境，安装好 torch、onnx、onnxsim 等第三方库

输出如下：

生成好的 resnet50.onnx 模型文件保存在 models/onnx 文件夹下，大家可以查看

接着我们需要利用 ONNX 生成对应的 engine 完成推理，在此之前我们需要修改下整体的 Makefile.config，指定一些库的路径：

# tensorrt_starter/config/Makefile.config
# CUDA_VER                    :=  11
CUDA_VER                    :=  11.6
    
# opencv和TensorRT的安装目录
OPENCV_INSTALL_DIR          :=  /usr/local/include/opencv4
# TENSORRT_INSTALL_DIR        :=  /mnt/packages/TensorRT-8.4.1.5
TENSORRT_INSTALL_DIR        :=  /home/jarvis/lean/TensorRT-8.6.1.6

Note：大家查看自己的 CUDA 是多少版本，修改为对应版本即可，另外 OpenCV 和 TensorRT 修改为你自己安装的路径即可

接着我们就可以来执行编译，指令如下：

make -j64

输出如下：

接着执行：

./bin/trt-infer

输出如下：

我们这里可以看到 INT8 模型在校准时的一些打印信息，它会读取校准图片通过校准器获取 INT8 量化时的一些 dynamic range 等信息，并将量化信息保存到校准文件中，所以我们在 calibration 文件夹中可以看到生成的 calibration_table.txt 校准文件，如下图所示，有了校准文件下次我们进行 PTQ 量化时只需要加载校准文件读取其中的量化信息进行量化即可

同时我们还详细打印了每层的信息，可以看到其中 weights 的 type 为 INT8，此外 INT8 量化后的 resenet50 模型在推理一张图片时仅需要 0.4ms 左右，我们对比看下 FP16 模型：

可以看到 FP16 模型推理部分耗时大约是 0.6ms，比 INT8 模型还是要慢的，不过其精度要高一些，我们从推理图片的置信度可以看出来

Note：这里大家可以测试下其它的 ONNX 模型看下推理结果，这里博主准备了导出好的各个分类模型的 ONNX，大家可以点击 here 下载

如果大家能够看到上述输出结果，那就说明本小节案例已经跑通，下面我们就来看看具体的代码实现

2. 补充说明

在分析代码之前我们先来看下韩君老师在这小节中写的 README 文档

这个小节主要介绍了如何用 C++ 实现 int8 的 calibrator，主要是以下几个文件：

• src/cpp/trt_calibrator.cpp
• src/cpp/trt_model.cpp
• include/trt-calibrator.hpp

如果我们想让模型实现 INT8 量化的话，需要在模型创建的时候在 config 里面设置，比如：

if (builder->platformHasFastFp16() && m_params->prec == model::FP16) {
    config->setFlag(BuilderFlag::kFP16);
    config->setFlag(BuilderFlag::kPREFER_PRECISION_CONSTRAINTS);
} else if (builder->platformHasFastInt8() && m_params->prec == model::INT8) {
    config->setFlag(BuilderFlag::kINT8);
    config->setFlag(BuilderFlag::kPREFER_PRECISION_CONSTRAINTS);
}

如果是 FP16 精度的话其实两行设置就可以了，但是如果是 INT8 的话，我们还需要自己设计一个 calibrator（校准器），这个和我们之前设计 logger 和 Plugin 一样，在创建 calibrator 类的时候需要继承 nvinfer1 里的 calibrator，NVIDIA 官方提供了以下五种校准器：

• nvinfer1::IInt8EntropyCalibrator2
• nvinfer1::IInt8MinMaxCalibrator
• nvinfer1::IInt8EntropyCalibrator
• nvinfer1::IInt8LegacyCalibrator
• nvinfer1::IInt8Calibrator

不同的 calibrator 所能够实现的 dynamic range 是不一样的，具体有什么区别大家感兴趣的可以看下：四. TensorRT模型部署优化-quantization(calibration)

我们在 calibrator 类中需要实现重载的函数有以下四个：

int         getBatchSize() const noexcept override {return m_batchSize;};
bool        getBatch(void* bindings[], const char* names[], int nbBindings) noexcept override;
const void* readCalibrationCache(std::size_t &length) noexcept override;
void        writeCalibrationCache (const void* ptr, std::size_t legth) noexcept override;

• getBatchSize：获取 calibration 的 batch 大小，需要注意的是不同的 batch size 会有不同的校准效果
• getBatch：calibration 是获取一个 batch 的图像，做完预处理之后 H2D 到 GPU，在 GPU 上做统计。这里需要注意的是在 getBatch 获取的图像必须和推理时所采用的预处理保持一致，不然 dynamic range 会不准
• readCalibrationCache：用来读取 calibration table，也就是之前做 calibration 统计得到的各个 layer 输出 tensor 的 dynamic range。实现这个函数可以让我们避免每次做 INT8 推理的时候都需要做一次 calibration
• writeCalibrationCache：将统计得到的 dynamic range 写入到 calibration table 中去

实现完了基本的 calibrator 之后，在 build 引擎的时候通过 config 指定 calibrator 就可以了，如下所示：

shared_ptr<Int8EntropyCalibrator> calibrator(new Int8EntropyCalibrator(
    64, 
    "calibration/calibration_list_imagenet.txt", 
    "calibration/calibration_table.txt",
    3 * 224 * 224, 224, 224));
config->setInt8Calibrator(calibrator.get());

这里面的 calibration_list_imagenet.txt 使用的是 ImageNet2012 的 test 数据集中的一部分，这里可以根据自己的情况更改。需要注意的是如果 calibrator 改变了，或者模型架构改变了，需要删除掉 calibration_table.txt 来重新计算 dynamic range，否则会报错。

下面我们来看具体的代码分析

3. 代码分析

3.1 main.cpp

我们先从 main.cpp 看起：

#include "trt_model.hpp"
#include "trt_logger.hpp"
#include "trt_worker.hpp"
#include "utils.hpp"

using namespace std;

int main(int argc, char const *argv[])
{
    /*这么实现目的在于让调用的整个过程精简化*/
    string onnxPath    = "models/onnx/resnet50.onnx";

    auto level         = logger::Level::VERB;
    auto params        = model::Params();

    params.img         = {224, 224, 3};
    params.num_cls     = 1000;
    params.task        = model::task_type::CLASSIFICATION;
    params.dev         = model::device::GPU;
    params.prec        = model::precision::INT8;

    // 创建一个worker的实例, 在创建的时候就完成初始化
    auto worker   = thread::create_worker(onnxPath, level, params);

    // 根据worker中的task类型进行推理
    worker->inference("data/cat.png");
    worker->inference("data/gazelle.png");
    worker->inference("data/eagle.png");
    worker->inference("data/fox.png");
    worker->inference("data/tiny-cat.png");
    worker->inference("data/wolf.png");

    return 0;
}

相比于之前的代码这里多了一个 INT8 精度的指定，表示我们想生成 INT8 量化后的模型，其它的部分和前面没有什么区别

大部分代码和上小节案例一样，所以这里博主就不再分析了，我们主要来看 INT8 量化的部分

3.2 trt_model.cpp

在 build_engine 函数中我们有所修改：

bool Model::build_engine() {
    // 我们也希望在build一个engine的时候就把一系列初始化全部做完，其中包括
    //  1. build一个engine
    //  2. 创建一个context
    //  3. 创建推理所用的stream
    //  4. 创建推理所需要的device空间
    // 这样，我们就可以在build结束以后，就可以直接推理了。这样的写法会比较干净
    auto builder       = shared_ptr<IBuilder>(createInferBuilder(*m_logger), destroy_trt_ptr<IBuilder>);
    auto network       = shared_ptr<INetworkDefinition>(builder->createNetworkV2(1), destroy_trt_ptr<INetworkDefinition>);
    auto config        = shared_ptr<IBuilderConfig>(builder->createBuilderConfig(), destroy_trt_ptr<IBuilderConfig>);
    auto parser        = shared_ptr<IParser>(createParser(*network, *m_logger), destroy_trt_ptr<IParser>);

    config->setMaxWorkspaceSize(m_workspaceSize);
    config->setProfilingVerbosity(ProfilingVerbosity::kDETAILED); //这里也可以设置为kDETAIL;

    if (!parser->parseFromFile(m_onnxPath.c_str(), 1)){
        return false;
    }

    shared_ptr<Int8EntropyCalibrator> calibrator;

    if (builder->platformHasFastFp16() && m_params->prec == model::FP16) {
        config->setFlag(BuilderFlag::kFP16);
        config->setFlag(BuilderFlag::kPREFER_PRECISION_CONSTRAINTS);
    } else if (builder->platformHasFastInt8() && m_params->prec == model::INT8) {
        config->setFlag(BuilderFlag::kINT8);
        config->setFlag(BuilderFlag::kPREFER_PRECISION_CONSTRAINTS);
        /* 
         * 新添加的针对int8 calibration做的操作 
         * 注意: 这里calibration的batch size会影响int8的精度
         */
        calibrator = shared_ptr<Int8EntropyCalibrator>(new Int8EntropyCalibrator(
            64, 
            "calibration/calibration_list_imagenet.txt", 
            "calibration/calibration_table.txt",
            3 * 224 * 224, 224, 224));
        config->setInt8Calibrator(calibrator.get());
    }


    auto engine        = shared_ptr<ICudaEngine>(builder->buildEngineWithConfig(*network, *config), destroy_trt_ptr<ICudaEngine>);
    auto plan          = builder->buildSerializedNetwork(*network, *config);
    auto runtime       = shared_ptr<IRuntime>(createInferRuntime(*m_logger), destroy_trt_ptr<IRuntime>);

    // 保存序列化后的engine
    save_plan(*plan);

    // 根据runtime初始化engine, context, 以及memory
    setup(plan->data(), plan->size());

    // 把优化前和优化后的各个层的信息打印出来
    LOGV("Before TensorRT optimization");
    print_network(*network, false);
    LOGV("After TensorRT optimization");
    print_network(*network, true);

    return true;
}

首先如果设备支持 INT8 量化并且 Model 参数的精度设置的是 INT8，那我们我们需要在 config 中设置 INT8 的标志如下所示：

else if (builder->platformHasFastInt8() && m_params->prec == model::INT8) {
    config->setFlag(BuilderFlag::kINT8);
    config->setFlag(BuilderFlag::kPREFER_PRECISION_CONSTRAINTS);
    ...
}

接着我们创建了一个校准器，并在 config 中进行了设置：

calibrator = shared_ptr<Int8EntropyCalibrator>(new Int8EntropyCalibrator(
    64, 
    "calibration/calibration_list_imagenet.txt", 
    "calibration/calibration_table.txt",
    3 * 224 * 224, 224, 224));
config->setInt8Calibrator(calibrator.get());

Int8EntropyCalibrator 是我们自定义实现的一个 calibrator 类，它继承自 nvinfer1::IInt8EntropyCalibrator2，它的构造函数的参数如下：

• batchSize：校准过程中每次出来的批量大小
• calibrationSetPath：校准数据集文件的路径，这个一个 .txt 文本文件，里面保存了校准图片的路径
• calibrationTablePath：校准表的保存路径，它记录了模型中每个 layer 的激活值分布以及相应的量化缩放因子
• inputSize：输入张量的大小，即 c*h*w 表示没模型输入的总大小
• inputH：输入图像的高度
• inputW：输入图像的宽度

config->setInt8Calibrator(calibrator.get()) 将校准器设置到 engine 配置中，以便在构建 INT8 engine 时使用

3.3 trt_calibrator.hpp

下面我们来看下自定义的 calibrator 类具体是如何实现的，我们先从头文件看起：

#ifndef __TRT_CALIBRATOR_HPP__
#define __TRT_CALIBRATOR_HPP__

#include "NvInfer.h"
#include <string>
#include <vector>


namespace model{
/*
 * 自定义一个calibrator类
 * 我们在创建calibrator的时候需要继承nvinfer1中的calibrator类
 * TensorRT提供了五种Calibrator类
 *
 *   - nvinfer1::IInt8EntropyCalibrator2
 *   - nvinfer1::IInt8MinMaxCalibrator
 *   - nvinfer1::IInt8EntropyCalibrator
 *   - nvinfer1::IInt8LegacyCalibrator
 *   - nvinfer1::IInt8Calibrator
 * 具体有什么不同，建议读一下官方文档和回顾一下之前的学习资料
 *
 * 默认下推荐使用IInt8EntropyCalibrator2
*/

class Int8EntropyCalibrator: public nvinfer1::IInt8EntropyCalibrator2 {
// class Int8EntropyCalibrator: public nvinfer1::IInt8MinMaxCalibrator {

public:
    Int8EntropyCalibrator(
        const int& batchSize,
        const std::string& calibrationSetPath,
        const std::string& calibrationTablePath,
        const int& inputSize,
        const int& inputH,
        const int& inputW);

    ~Int8EntropyCalibrator(){};

    int         getBatchSize() const noexcept override {return m_batchSize;};
    bool        getBatch(void* bindings[], const char* names[], int nbBindings) noexcept override;
    const void* readCalibrationCache(std::size_t &length) noexcept override;
    void        writeCalibrationCache (const void* ptr, std::size_t legth) noexcept override;

private:
    const int   m_batchSize;
    const int   m_inputH;
    const int   m_inputW;
    const int   m_inputSize;
    const int   m_inputCount;
    const std::string m_calibrationTablePath {nullptr};
    
    std::vector<std::string> m_imageList;
    std::vector<char>        m_calibrationCache;

    float* m_deviceInput{nullptr};
    bool   m_readCache{true};
    int    m_imageIndex;
    
};

}; // namespace model

#endif __TRT_CALIBRATOR_HPP__

这个自定义的 Int8EntropyCalibrator 类头文件定义了一个用于 TensorRT INT8 量化的校准器类。该类继承自 nvinfer1::IInt8EntropyCalibrator2，其构造函数我们前面已经分析过了，重载的成员函数就是我们之前提到的四个：

• getBatchSize：TensorRT 在校准过程中会调用这个函数来确定每次出来多少张图片
• getBatch：TensorRT 在校准时会调用这个函数来获取一个 batch 的数据，并将数据绑定到相应的输入
• readCalibrationCache：如果校准表已经存在，可以直接读取以加快校准过程，而不必重新执行完整的校准
• writeCalibrationCache：在校准过程结束后，将生成的校准表保存到指定路径，以供后续使用

成员变量主要有以下几个：

• m_batchSize：批量大小
• m_inputH 和 m_inputW：输入图像的高度和宽度
• m_inputSize：输入图片张量的总大小
• m_inputCount：batch 张图片张量的总大小
• m_calibrationTablePath：校准表的保存路径
• m_imageList：存储所有待校准的图像文件路径
• m_calibrationCache：用于存储从缓存中读取或将要写入缓存的校准数据
• m_deviceInput：指向 device 输入内存的指针，用于在 device 上存储输入数据，以便进行校准
• m_readCache：标记是否需要读取缓存数据
• m_imageIndex：当前处理的图像索引，用于追踪校准过程中处理到的数据位置

这个自定义的 Int8EntropyCalibrator 类提供了一种对 TensorRT 的 INT8 量化校准进行定制的方法。通过继承 IInt8EntropyCalibrator2 并实现其接口函数，用户可以加载自定义的数据集进行校准，并利用缓存机制来加快校准过程

3.4 trt_calibrator.cpp

下面我们来看这个类具体的实现：

#include "NvInfer.h"
#include "trt_calibrator.hpp"
#include "utils.hpp"
#include "trt_logger.hpp"
#include "trt_preprocess.hpp"

#include <fstream>
#include <vector>
#include <algorithm>
#include <iterator>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>

using namespace std;
using namespace nvinfer1;

namespace model{

/*
 * calibrator的构造函数
 * 我们在这里把calibration所需要的数据集准备好，需要保证数据集的数量可以被batchSize整除
 * 同时由于calibration是在device上进行的，所以需要分配空间
 */
Int8EntropyCalibrator::Int8EntropyCalibrator(
    const int&    batchSize,
    const string& calibrationDataPath,
    const string& calibrationTablePath,
    const int&    inputSize,
    const int&    inputH,
    const int&    inputW):

    m_batchSize(batchSize),
    m_inputH(inputH),
    m_inputW(inputW),
    m_inputSize(inputSize),
    m_inputCount(batchSize * inputSize),
    m_calibrationTablePath(calibrationTablePath)
{
    m_imageList = loadDataList(calibrationDataPath);
    m_imageList.resize(static_cast<int>(m_imageList.size() / m_batchSize) * m_batchSize);
    std::random_shuffle(m_imageList.begin(), m_imageList.end(), 
                        [](int i){ return rand() % i; });
    CUDA_CHECK(cudaMalloc(&m_deviceInput, m_inputCount * sizeof(float)));
}

/*
 * 获取做calibration的时候的一个batch的图片，之后上传到device上
 * 需要注意的是，这里面的一个batch中的每一个图片，都需要做与真正推理是一样的前处理
 * 这里面我们选择在GPU上进行前处理，所以处理万
 */
bool Int8EntropyCalibrator::getBatch(
    void* bindings[], const char* names[], int nbBindings) noexcept
{
    if (m_imageIndex + m_batchSize >= m_imageList.size() + 1)
        return false;

    LOG("%3d/%3d (%3dx%3d): %s", 
        m_imageIndex + 1, m_imageList.size(), m_inputH, m_inputW, m_imageList.at(m_imageIndex).c_str());
    
    /*
     * 对一个batch里的所有图像进行预处理
     * 这里可有以及个扩展的点
     *  1. 可以把这个部分做成函数，以函数指针的方式传给calibrator。因为不同的task会有不同的预处理
     *  2. 可以实现一个bacthed preprocess
     * 这里留给当作今后的TODO
     */
    cv::Mat input_image;
    float mean[]       = {0.406, 0.456, 0.485};
    float std[]        = {0.225, 0.224, 0.229};
    for (int i = 0; i < m_batchSize; i ++){
        input_image = cv::imread(m_imageList.at(m_imageIndex++));
        process::preprocess_resize_gpu(
            input_image, 
            m_deviceInput + i * m_inputSize,
            m_inputH, m_inputW, 
            mean, std, process::tactics::GPU_BILINEAR);
    }

    bindings[0] = m_deviceInput;

    return true;
}
    
/* 
 * 读取calibration table的信息来创建INT8的推理引擎, 
 * 将calibration table的信息存储到calibration cache，这样可以防止每次创建int推理引擎的时候都需要跑一次calibration
 * 如果没有calibration table的话就会直接跳过这一步，之后调用writeCalibrationCache来创建calibration table
 */
const void* Int8EntropyCalibrator::readCalibrationCache(size_t& length) noexcept
{
    void* output;
    m_calibrationCache.clear();

    ifstream input(m_calibrationTablePath, ios::binary);
    input >> noskipws;
    if (m_readCache && input.good())
        copy(istream_iterator<char>(input), istream_iterator<char>(), back_inserter(m_calibrationCache));

    length = m_calibrationCache.size();
    if (length){
        LOG("Using cached calibration table to build INT8 trt engine...");
        output = &m_calibrationCache[0];
    }else{
        LOG("Creating new calibration table to build INT8 trt engine...");
        output = nullptr;
    }
    return output;
}

/* 
 * 将calibration cache的信息写入到calibration table中
*/
void Int8EntropyCalibrator::writeCalibrationCache(const void* cache, size_t length) noexcept
{
    ofstream output(m_calibrationTablePath, ios::binary);
    output.write(reinterpret_cast<const char*>(cache), length);
    output.close();
}

} // namespace model

首先是构造函数部分：

Int8EntropyCalibrator::Int8EntropyCalibrator(
    const int&    batchSize,
    const string& calibrationDataPath,
    const string& calibrationTablePath,
    const int&    inputSize,
    const int&    inputH,
    const int&    inputW):

    m_batchSize(batchSize),
    m_inputH(inputH),
    m_inputW(inputW),
    m_inputSize(inputSize),
    m_inputCount(batchSize * inputSize),
    m_calibrationTablePath(calibrationTablePath)
{
    m_imageList = loadDataList(calibrationDataPath);
    m_imageList.resize(static_cast<int>(m_imageList.size() / m_batchSize) * m_batchSize);
    std::random_shuffle(m_imageList.begin(), m_imageList.end(), 
                        [](int i){ return rand() % i; });
    CUDA_CHECK(cudaMalloc(&m_deviceInput, m_inputCount * sizeof(float)));
}

在构造函数中我们进行了一些参数的初始化，并从指定的校准数据集路径中加载了所有图像文件路径，然后将图像列表调整为批量大小的整数倍，并随机打乱顺序，确保数据分布均匀，避免顺序带来的偏差

最后我们通过 cudaMalloc 在 GPU 上分配了一个足够的内存来存储一个 batch 的输入数据，其中 m_deviceInput 是一个指向 device 内存的指针

接着是 getBatch 函数：

bool Int8EntropyCalibrator::getBatch(
    void* bindings[], const char* names[], int nbBindings) noexcept
{
    if (m_imageIndex + m_batchSize >= m_imageList.size() + 1)
        return false;

    LOG("%3d/%3d (%3dx%3d): %s", 
        m_imageIndex + 1, m_imageList.size(), m_inputH, m_inputW, m_imageList.at(m_imageIndex).c_str());
    
    cv::Mat input_image;
    float mean[]       = {0.406, 0.456, 0.485};
    float std[]        = {0.225, 0.224, 0.229};
    for (int i = 0; i < m_batchSize; i ++){
        input_image = cv::imread(m_imageList.at(m_imageIndex++));
        process::preprocess_resize_gpu(
            input_image, 
            m_deviceInput + i * m_inputSize,
            m_inputH, m_inputW, 
            mean, std, process::tactics::GPU_BILINEAR);
    }

    bindings[0] = m_deviceInput;

    return true;
}

该函数的目的是获取一个 batch 的图像进行预处理后传输到 device 上，首先检查是否有足够的图像来处理当前 batch，如果没有则返回。接着循环当前 batch 的图像，调用 preprocess_resize_gpu 进行图像预处理，最后将预处理后的图像数据绑定到 TensorRT 的输入绑定点，以便后续校准过程使用

值得注意的是这里有个扩展点，我们可以把上述这个预处理部分做成一个函数，以函数指针的方式传给 calibrator，这是因为不同的 task 会有不同的预处理，我们最好不要写死，实现一个 batched preprocess 就行，在 tensorRT_Pro 中的 int8process 就是这么做的，大家感兴趣的可以看看：app_yolo.cpp#L121

然后是 readCalibrationCache 函数：

const void* Int8EntropyCalibrator::readCalibrationCache(size_t& length) noexcept
{
    void* output;
    m_calibrationCache.clear();

    ifstream input(m_calibrationTablePath, ios::binary);
    input >> noskipws;
    if (m_readCache && input.good())
        copy(istream_iterator<char>(input), istream_iterator<char>(), back_inserter(m_calibrationCache));

    length = m_calibrationCache.size();
    if (length){
        LOG("Using cached calibration table to build INT8 trt engine...");
        output = &m_calibrationCache[0];
    }else{
        LOG("Creating new calibration table to build INT8 trt engine...");
        output = nullptr;
    }
    return output;
}

该函数的功能是从校准表文件中读取已经存在的校准缓存，以加快 INT8 量化的校准过程，首先清空当前的校准缓存，以便重新加载，接着打开校准表文件并以二进制方式读取。如果允许读取缓存且文件打开成功，则从文件中读取数据并存储在 m_calibrationCache 中，如果读取了缓存数据则返回缓存的起始地址，否则返回 nullptr 表示没有可用的缓存

最后我们来看下 writeCalibrationCache 函数的实现：

void Int8EntropyCalibrator::writeCalibrationCache(const void* cache, size_t length) noexcept
{
    ofstream output(m_calibrationTablePath, ios::binary);
    output.write(reinterpret_cast<const char*>(cache), length);
    output.close();
}

该函数的功能主要是将校准缓存数据写入到校准表文件中，以便下次使用时可以直接加载缓存，避免重复校准

OK，以上就是自定义校准器 calibrator 类的具体实现

4. 校准精度影响因素

我们来看下不同的 batch size 对精度的影响：

batch=1

batch=64

当然我们并没有一个测试集去验证这个模型的 top1 或者 top5 的性能，只是简单的通过一些图片的置信度去猜测，这个其实是不准的，大家感兴趣的可以自己准备一个数据集测试看看校准时 batch size 的大小是否会影响最终的精度

另外大家可以测试下不同的 calibrator，修改也很简单，只需要在 trt_calibrator.hpp 中继承不同的 calibrator 就行，示例如下：

// class Int8EntropyCalibrator: public nvinfer1::IInt8EntropyCalibrator2 {
class Int8EntropyCalibrator: public nvinfer1::IInt8MinMaxCalibrator {

最后大家还可以测试下不同数量的校准图片对模型的精度是否存在影响，选 500 张，2000 张的精度是否不同，是不是校准图片数量越多越好呢？

其实之前博主有做过一些简单的测试，大家感兴趣的可以看看：YOLOv5-PTQ量化部署 以及 YOLOv7-PTQ量化部署

Tips：大家如果在 Jetson 嵌入式端创建 engine 时觉得时间太长的话可以使用 timing caches 来加速

TensorRT 的 timing caches 是一个功能，它可以加快 engine 的创建过程。这个特性主要是为了优化和加速内核选择过程。在创建 TensorRT engine 时，需要进行大量的内核调优（kernel tuning）和选择最佳的算法，这个过程通常会花费很多时间，尤其是在算力有限的设备如 NVIDIA Jetson 嵌入式设备上。

以下是关于 timing cache 的一些说明：(from ChatGPT)

Timing Cache 的工作原理：

• 1. 缓存生成： 在 engine 的编译过程中，TensorRT 会测试不同算法的性能，并将最优算法的性能数据存储在 timing cache 中。
• 2. 复用缓存： 当再次构建相似的 engine 时，TensorRT 可以直接查阅之前存储的 timing cache，而不需要重新进行性能测试。这大大减少了 engine 创建的时间。

如何使用 Timing Cache：

• 创建和保存 Timing Cache： 在第一次构建 engine 时，TensorRT 会创建一个 timing cache。你可以将这个 timing cache 保存到文件中，以便未来复用。
• 加载 Timing Cache： 在构建新的 engine 时，可以加载已经保存的 timing cache，从而避免重新进行时间消耗的内核调优过程。

结语

本次课程我们学习了分类模型的 PTQ 量化，自定义实现了一个 calibrator 校准器，主要是继承 nvinfer1 中的校准器，接着重载实现四个函数，总体来说还是比较简单的，当然 PTQ 量化中的一些讨论还是比较丰富的，例如 batch size 的大小，校准图片的数量以及校准器的选择等等

OK，以上就是 6.3 小节案例的全部内容了，下节我们来学习 6.4 小节的 trt-engine-explorer 工具包，敬请期待😄

下载链接

• tensorrt_starter源码
• 6.3-int8-calibration案例文件

参考

• Ubuntu20.04软件安装大全
• https://github.com/kalfazed/tensorrt_starter
• https://www.kaggle.com/datasets?search=ILSVRC2012
• 四. TensorRT模型部署优化-quantization(calibration)