模型量化!ONNX转TensorRT(FP32, FP16, INT8)

这篇具有很好参考价值的文章主要介绍了模型量化!ONNX转TensorRT(FP32, FP16, INT8)。希望对大家有所帮助。如果存在错误或未考虑完全的地方,请大家不吝赐教,您也可以点击"举报违法"按钮提交疑问。

本文为Python实现,C++实现链接模型量化
若还没有配置环境(CUDA,CUDNN,TensorRT),请移至C++实现中查看环境配置方法
  
支持三种不同精度的量化

模型单精度量化 (FP32)
模型半精度量化 (FP16)
模型Int8量化 (INT8)

经测试yolov5,yolov6,yolov7,yolov8转化成功
yolov5: https://github.com/ultralytics/yolov5
yolov6: https://github.com/meituan/YOLOv6
yolov7: https://github.com/WongKinYiu/yolov7
yolov8: https://github.com/ultralytics/ultralytics

注意:若使用INT8量化,需要额外的文件,这里是calibrator.py

main.py文件
import tensorrt as trt
import os
from calibrator import Calibrator

LOGGER = trt.Logger(trt.Logger.VERBOSE)

def buildEngine(onnx_file, engine_file, quantification, batch_size, FP16_mode, INT8_mode,
                img_height, img_wdith, calibration_images, calibration_cache):
    builder = trt.Builder(LOGGER)
    network = builder.create_network(1<<int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, LOGGER)
    config = builder.create_builder_config()
    parser.parse_from_file(onnx_file)

    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 16*(1<<20))

    if FP16_mode == True:
        config.set_flag(trt.BuilderFlag.FP16)
    elif INT8_mode == True:
        config.set_flag(trt.BuilderFlag.INT8)
        config.int8_calibrator = Calibrator(quantification, batch_size, img_height, img_wdith, calibration_images, calibration_cache)
    engine = builder.build_serialized_network(network, config)
    if engine is None:
        print("EXPORT ENGINE FAILED!")
        # exit(0)

    with open(engine_file, "wb") as f:
        f.write(engine)


def main():
    quantification = 1  # 量化次数
    batch_size = 1
    img_height = 640
    img_wdith = 640
    calibration_images = "./images"
    onnx_file = "./weights/yolov7.onnx"
    engine_file = "./weights/yolov7.engine"
    calibration_cache = "./weights/yolov7_calibration.cache"

    '''
    模型使用单精度量化,设置 FP16_mode = False & INT8_mode = False  (FP32)
    模型使用半精度量化,设置 FP16_mode = True  & INT8_mode = False  (FP16)
    模型使用 Int8量化,设置 FP16_mode = False & INT8_mode = True   (INT8)
    '''
    FP16_mode = False
    INT8_mode = False

    if not os.path.exists(onnx_file):
        print("LOAD ONNX FILE FAILED: ", onnx_file)

    print('Load ONNX file from:%s \nStart export, Please wait a moment...'%(onnx_file))
    buildEngine(onnx_file, engine_file, quantification, batch_size, FP16_mode, INT8_mode,
                img_height, img_wdith, calibration_images, calibration_cache)
    print('Export ENGINE success, Save as: ', engine_file)


if __name__ == '__main__':
    main()
calibrator.py文件
import os
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import logging
import numpy as np
import cv2

class Calibrator(trt.IInt8EntropyCalibrator):
    def __init__(self, quantification=1, batch_size=1, height=640, width=640, calibration_images="", cache_file=""):
        trt.IInt8EntropyCalibrator.__init__(self)
        self.index = 0
        self.length = quantification
        self.batch_size = batch_size
        self.cache_file = cache_file
        self.height = height
        self.width = width
        self.img_list = [calibration_images + '/' + l for l in os.listdir(calibration_images)]
        self.calibration_data = np.zeros((self.batch_size, 3, self.height, self.width), dtype=np.float32)
        self.d_input = cuda.mem_alloc(self.calibration_data.nbytes)

    def next_batch(self):
        if self.index < self.length:
            for i in range(self.batch_size):
                img = cv2.imread(self.img_list[i + self.index*self.batch_size])
                img = self.preprocess(img)
                self.calibration_data[i] = img
            self.index += 1
            return np.ascontiguousarray(self.calibration_data, dtype=np.float32)
        else:
            return np.array([])

    def __len__(self):
        return self.length

    def get_batch_size(self):
        return self.batch_size

    def get_batch(self, name):
        batch = self.next_batch()
        if not batch.size:
            return None
        cuda.memcpy_htod(self.d_input, batch)
        return [int(self.d_input)]

    def read_calibration_cache(self):
        # If there is a cache, use it instead of calibrating again. Otherwise, implicitly return None.
        if os.path.exists(self.cache_file):
            with open(self.cache_file, "rb") as f:
                return f.read()

    def write_calibration_cache(self, cache):
        with open(self.cache_file, "wb") as f:
            f.write(cache)

    def preprocess(self, img):
        h, w, c = img.shape
        image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        r_w = self.width / w
        r_h = self.height / h
        if r_h > r_w:
            tw = self.width
            th = int(r_w * h)
            tx1 = tx2 = 0
            ty1 = int((self.height - th) / 2)
            ty2 = self.height - th - ty1
        else:
            tw = int(r_h * w)
            th = self.height
            tx1 = int((self.width - tw) / 2)
            tx2 = self.width - tw - tx1
            ty1 = ty2 = 0
        image = cv2.resize(image, (tw, th))
        image = cv2.copyMakeBorder(image, ty1, ty2, tx1, tx2, cv2.BORDER_CONSTANT, (128, 128, 128))
        image = image.astype(np.float32)
        image /= 255.0
        image = np.transpose(image, [2, 0, 1])
        return image

结束~文章来源地址https://www.toymoban.com/news/detail-604342.html

到了这里,关于模型量化!ONNX转TensorRT(FP32, FP16, INT8)的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处: 如若内容造成侵权/违法违规/事实不符,请点击违法举报进行投诉反馈,一经查实,立即删除!

领支付宝红包 赞助服务器费用

相关文章

  • 【深度学习】SDXL tensorRT 推理,Stable Diffusion 转onnx,转TensorRT

    juggernautXL_version6Rundiffusion.safetensors文件是pth pytroch文件,需要先转为diffusers 的文件结构。 FP16在后面不好操作,所以最好先是FP32: 有了diffusers 的文件结构,就可以转onnx文件。 项目:https://huggingface.co/docs/diffusers/optimization/onnx stabilityai/stable-diffusion-xl-1.0-tensorrt 项目:https://hug

    2024年01月19日
    浏览(54)
  • onnx转TensorRT使用的三种方式(最终在Python运行)

    背景 记录下onnx转成TensorRT加速的三种方式 1. 直接使用onnxruntime 在onnxruntime的session初始化的时候第一个provider加入TensorrtExecutionProvider,软件会自动查找是否支持TensorRT,如果可以就会进行转换并运行,如果不可以会接着找下一个,也有可能TensorRT跑一半报错,这就得看环境什么

    2024年02月08日
    浏览(42)
  • YOLOV5使用(一): docker跑通,详解TensorRT下plugin的onnx

    使用ubuntu为案例 使用端口映射功能也就是说打开jupyter lab的指令是 当然,个人建议直接去vscode端口点击就打开jupyter lab和tensorboard比较方便 YOLO格式的标签文件是一个纯文本文件,每个文件名对应一张图像,每个标签文件中包含了该图像中所有检测到的目标的信息。 YOLOv5的标

    2024年02月04日
    浏览(43)
  • 【AI】yolov5:数据下载、图像标注、训练、预测等,超详细笔记;pytorch转onnx再转TensorRT;训练识别口罩

    https://gitcode.net/mirrors/ultralytics/yolov5/ 1)安装pytorch 2)安装cudnn 3)下载yolov5源码: 4)安装yolov5依赖

    2024年02月15日
    浏览(44)
  • AI大模型的神经网络模型量化技术:INT8 还是 INT4 ?

    目录 背景介绍 模型量化 常见技术 量化工具 量化精度的选择 结束语

    2024年02月13日
    浏览(35)
  • 【MMDeploy&MMPose】手把手教你在Windows上使用MMDeploy1.x进行ONNX和TensorRT的部署(Python篇)

    2023.9.8更新 C++ SDK篇已经发布,点击这里传送到文章 MMDeploy 提供了一系列工具,帮助我们更轻松的将 OpenMMLab 下的算法部署到各种设备与平台上。 MMDeploy 定义的模型部署流程,这里直接使用了官方文档的内容,如下图所示:  1)模型转换(Model Converter) 模型转换的主要功能是

    2024年01月16日
    浏览(46)
  • Golang中int, int8, int16, int32, int64和uint区别

    在Go语言(Golang)中,整数类型用来表示整数,主要包括int系列和uint系列。这两大系列根据能否表示负数,分为有符号整数和无符号整数。 int系列( 有符号整数 ): int8: 占用1个字节(8位),能够表示的有符号整数范围从 -128到127。 int16: 占用2个字节(16位),能够表示的有

    2024年01月25日
    浏览(58)
  • OpenVINO 2022.3实战六:NNCF 实现 YOLOv5 模型 INT8 量化

    使用OpenVINO模型优化器将YOLOv5模型转换为OpenVINO IR格式,以便在Intel硬件上进行推理。 下载yolov5代码 ultralytics/yolov5 导出模型为onnx模型,接着使用mo导出openvino fp32和fp16模型 将训练数据集准备成可用于量化的格式。 配置量化管道,例如选择适当的量化算法和设置目标精度。 在

    2024年02月08日
    浏览(41)
  • 实现CenterNet图像分割算法模型的转换和量化(SDK0301-转ONNX编译)

    一、实现CenterNet图像分割算法模型的转换和量化(SDK0301-转ONNX编译) 1、模型转换 (1)下载CenterNet算法移植代码: 获取主干网的预训练模型,运行如下命令下载 dlav0 作为主干网的预训练模型: (2)将pth文件转换为ONNX格式模型 【SDK 230501版本后支持Torch模型权重直接转换】

    2024年02月09日
    浏览(36)
  • 详解C语言中的int8_t、uint8_t、int16_t、uint16_t、int32_t、uint32_t、int64_t、uint64_t

    2023年8月8日,周二上午 目录 为什么会产生int8_t、uint8_t等这类数据类型 int8_t、uint8_t等这类数据类型有什么用 头文件 int8_t、uint8_t等这类数据类型是怎么实现的 根本原因在于,C 语言标准只是规定了各个基本类型的最小字节长度,但并没有指定它们的确切字节长度: char 类型

    2024年02月13日
    浏览(42)

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

博客赞助

微信扫一扫打赏

请作者喝杯咖啡吧~博客赞助

支付宝扫一扫领取红包,优惠每天领

二维码1

领取红包

二维码2

领红包