【yolov8目标检测部署】TensorRT int8量化

原作者github:https://github.com/xuanandsix/Tensorrt-int8-quantization-pipline/tree/main

改进:

源代码支持的TensorRT版本为7.许多属性已经弃用;

在原有的代码上将支持的TensorRT版本从7改到8.

!!不知道如何安装TensorRT?请参考:【通俗易懂】Windows系统安装TensorRT

!!项目排坑请参考:【Debug】TensorRT报错汇总 欢迎大家补充。

一、简介

模型量化是深度学习领域中一种重要的模型优化技术,它主要涉及将模型中的浮点数(如FP32)转换为低比特的格式(如INT8)。这样做的目的是为了减少模型的存储大小和内存占用,同时加快模型的推理速度,尤其在硬件资源受限的设备上效果显著。

模型量化的方法:

  1. 训练时量化也叫量化感知训练(Quantization-Aware-Training,QAT),在训练过程中模拟量化操作,使模型适应量化,减少量化误差。

  2. 训练后量化(Post-Training Quantization,PTQ),PTQ不需要再训练,因此是一种轻量级的量化方法。在模型训练好后,再通过一个校准(Calibration)流程去计算比例因子,从而实现量化过程。

这里使用第二种方法简单且高效。

二、量化流程

2.1 TensorRT提供了四种IInt8Calibrator:

  • Int8EntropyCalibrator2:当前推荐的熵校准器,默认情况下校准发生在层融合之前,推荐用于CNN模型中。
  • IInt8MinMaxCalibrator:该校准器使用激活分布的整个范围来确定比例因子,默认情况下校准发生在层融合之前,推荐用于NLP任务的模型中。
  • IInt8EntropyCalibrator: 该校准器是TensorRT最原始的熵校准器,默认情况下校准发生在层融合之后,目前已不推荐使用。
  • IInt8LegacyCalibrator: 该校准器需要用户进行参数化,默认情况下校准发生在层融合之后,不推荐使用。

2.2 IInt8Calibrator需要实现的功能

  • getBatchSize:获取校准过程中的batchSize;
  • getBatch:获取校准过程中的输入;
  • writeCalibrationCache:由于校准花费的时间比较长,调用该函数将校准参数结果写入本地文件,方便下次直接读取。
  • readCalibrationCache:读取保存在本地的校准参数文件,在生成引擎过程中会自动调用。.

2.3 TensorRT构建INT8模型引擎时,会执行下面的步骤:

构建一个32位的模型引擎,然后在校准数据集上运行这个引擎,然后为每个张量激活值的分布记录一个直方图;

从直方图构建一个校准表,为每个张量计算出一个比例因子;

根据校准表和模型的定义构建一个INT8的引擎。

校准的过程可能会比较慢,不过第二步生成的校准表可以输出到文件并可以被重用,如果校准表文件已存在,那么校准器就直接从该文件中读取校准表而无需执行前面两步。另外,与引擎文件不同的是,校准表是可以跨平台使用的。因此,我们在实际部署模型过程中可以先在带通用GPU的计算机上生成校准表,然后在Jetson Nano等嵌入式平台上去使用。为了编码方便,我们可以用Python编程来实现INT8量化过程来生成校准表。

三、关键代码:

加载校准文件:quantization.py

python 复制代码
import numpy as np
import util_trt
import glob,os,cv2

BATCH_SIZE = 4  # 批处理大小
BATCH = 100     # 校准批次
height = 640    # 输入高度
width = 640     # 输入宽度
CALIB_IMG_DIR = './calibration/'    # 校准图片路径
onnx_model_path = './weights/yolov8s.onnx'  # 待校准onnx模型路径

def preprocess(img):
    img = cv2.resize(img, (height, width))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = img.transpose((2, 0, 1)).astype(np.float32)
    img /= 255.0
    return img

class DataLoader:
    def __init__(self):
        self.index = 0
        self.length = BATCH
        self.batch_size = BATCH_SIZE
        # self.img_list = [i.strip() for i in open('calib.txt').readlines()]
        self.img_list = glob.glob(os.path.join(CALIB_IMG_DIR, "*.jpg"))
        assert len(self.img_list) > self.batch_size * self.length, '{} must contains more than '.format(CALIB_IMG_DIR) + str(self.batch_size * self.length) + ' images to calib'
        print('found all {} images to calib.'.format(len(self.img_list)))
        self.calibration_data = np.zeros((self.batch_size,3,height,width), dtype=np.float32)

    def reset(self):
        self.index = 0

    def next_batch(self):
        if self.index < self.length:
            for i in range(self.batch_size):
                assert os.path.exists(self.img_list[i + self.index * self.batch_size]), 'not found!!'
                img = cv2.imread(self.img_list[i + self.index * self.batch_size])
                img = preprocess(img)
                self.calibration_data[i] = img

            self.index += 1

            # example only
            return np.ascontiguousarray(self.calibration_data, dtype=np.float32)
        else:
            return np.array([])

    def __len__(self):
        return self.length

def main():
    # onnx2trt
    fp16_mode = False
    int8_mode = True 
    print('*** onnx to tensorrt begin ***')
    # calibration
    calibration_stream = DataLoader()
    engine_model_path = "./weights/yolov8s_Int8.trt"
    calibration_table = 'best_calibration.cache'
    # fixed_engine,校准产生校准表
    engine_fixed = util_trt.get_engine(BATCH_SIZE, onnx_model_path, engine_model_path, fp16_mode=fp16_mode, 
        int8_mode=int8_mode, calibration_stream=calibration_stream, calibration_table_path=calibration_table, save_engine=True)
    print(engine_fixed)
    assert engine_fixed, 'Brokenls engine_fixed'
    print('*** onnx to tensorrt completed ***\n')
    
if __name__ == '__main__':
    main()

四、测试结果

model(640*640) infer_time(ms)/frame size
float32(trt) 8. 51.6M
int8(trt) 2. 14.3M

yolov8s模型FP32和INT8精度的目标检测结果分别如下面两张图片所示:

float32 tensorrt int8 tensorrt
相关推荐
车载诊断技术37 分钟前
电子电气架构 --- 什么是EPS?
网络·人工智能·安全·架构·汽车·需求分析
KevinRay_42 分钟前
Python超能力:高级技巧让你的代码飞起来
网络·人工智能·python·lambda表达式·列表推导式·python高级技巧
跃跃欲试-迪之1 小时前
animatediff 模型网盘分享
人工智能·stable diffusion
Captain823Jack1 小时前
nlp新词发现——浅析 TF·IDF
人工智能·python·深度学习·神经网络·算法·自然语言处理
被制作时长两年半的个人练习生1 小时前
【AscendC】ReduceSum中指定workLocal大小时如何计算
人工智能·算子开发·ascendc
Captain823Jack2 小时前
w04_nlp大模型训练·中文分词
人工智能·python·深度学习·神经网络·算法·自然语言处理·中文分词
Black_mario2 小时前
链原生 Web3 AI 网络 Chainbase 推出 AVS 主网, 拓展 EigenLayer AVS 应用场景
网络·人工智能·web3
Aileen_0v02 小时前
【AI驱动的数据结构:包装类的艺术与科学】
linux·数据结构·人工智能·笔记·网络协议·tcp/ip·whisper
数信云 DCloud2 小时前
实力认可 | 通付盾入选《ISC.AI 2024创新能力全景图谱》五项领域
人工智能
itwangyang5202 小时前
AIDD - 从机器学习到深度学习:蛋白质-配体对接评分函数的进展
人工智能·深度学习·机器学习