【yolov5】onnx的INT8量化engine

news2024/11/19 21:20:19

GitHub上有大佬写好代码,理论上直接克隆仓库里下来使用

git clone https://github.com/Wulingtian/yolov5_tensorrt_int8_tools.git

然后在yolov5_tensorrt_int8_tools的convert_trt_quant.py 修改如下参数

BATCH_SIZE 模型量化一次输入多少张图片

BATCH 模型量化次数

height width 输入图片宽和高

CALIB_IMG_DIR 训练图片路径,用于量化

onnx_model_path onnx模型路径

engine_model_path 模型保存路径

其中这个batch_size不能超过照片的数量,然后跑这个convert_trt_quant.py

出问题了吧@_@

这是因为tensor的版本更新原因,这个代码的tensorrt版本是7系列的,而目前新的tensorrt版本已经没有了一些属性,所以我们需要对这个大佬写的代码进行一些修改

如何修改呢,其实tensorrt官方给出了一个caffe量化INT8的例子

https://github.com/NVIDIA/TensorRT/tree/master/samples/python/int8_caffe_mnist

如果足够NB是可以根据官方的这个例子修改一下直接实现onnx的INT8量化的

但是奈何我连半桶水都没有,只有一滴水,但是这个例子中的tensorrt版本是新的,于是我尝试将上面那位大佬的代码修改为使用新版的tensorrt

居然成功了??!!

成功量化后的模型大小只有4MB,相比之下的FP16的大小为6MB,FP32的大小为9MB

再看看检测速度,速度和FP16差不太多

但是效果要差上一些了

那肯定不能忘记送上修改的代码,折腾一晚上的结果如下,主要是 util_trt程序

# tensorrt-lib

import os
import tensorrt as trt
import pycuda.autoinit
import pycuda.driver as cuda
from calibrator import Calibrator
from torch.autograd import Variable
import torch
import numpy as np
import time
# add verbose
TRT_LOGGER = trt.Logger(trt.Logger.VERBOSE) # ** engine可视化 **

# create tensorrt-engine
  # fixed and dynamic
def get_engine(max_batch_size=1, onnx_file_path="", engine_file_path="",\
               fp16_mode=False, int8_mode=False, calibration_stream=None, calibration_table_path="", save_engine=False):
    """Attempts to load a serialized engine if available, otherwise builds a new TensorRT engine and saves it."""
    def build_engine(max_batch_size, save_engine):
        """Takes an ONNX file and creates a TensorRT engine to run inference with"""
        with trt.Builder(TRT_LOGGER) as builder, \
                builder.create_network(1) as network,\
                trt.OnnxParser(network, TRT_LOGGER) as parser:
            
            # parse onnx model file
            if not os.path.exists(onnx_file_path):
                quit('ONNX file {} not found'.format(onnx_file_path))
            print('Loading ONNX file from path {}...'.format(onnx_file_path))
            with open(onnx_file_path, 'rb') as model:
                print('Beginning ONNX file parsing')
                parser.parse(model.read())
                assert network.num_layers > 0, 'Failed to parse ONNX model. \
                            Please check if the ONNX model is compatible '
            print('Completed parsing of ONNX file')
            print('Building an engine from file {}; this may take a while...'.format(onnx_file_path))        
            
            # build trt engine
            builder.max_batch_size = max_batch_size
            config = builder.create_builder_config()
            config.max_workspace_size = 1 << 20
            if int8_mode:
                config.set_flag(trt.BuilderFlag.INT8)
                assert calibration_stream, 'Error: a calibration_stream should be provided for int8 mode'
                config.int8_calibrator  = Calibrator(calibration_stream, calibration_table_path)
                print('Int8 mode enabled')
            runtime=trt.Runtime(TRT_LOGGER)
            plan = builder.build_serialized_network(network, config)
            engine = runtime.deserialize_cuda_engine(plan)
            if engine is None:
                print('Failed to create the engine')
                return None   
            print("Completed creating the engine")
            if save_engine:
                with open(engine_file_path, "wb") as f:
                    f.write(engine.serialize())
            return engine
        
    if os.path.exists(engine_file_path):
        # If a serialized engine exists, load it instead of building a new one.
        print("Reading engine from file {}".format(engine_file_path))
        with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
            return runtime.deserialize_cuda_engine(f.read())
    else:
        return build_engine(max_batch_size, save_engine)

唔,convert_trt_quant.py的代码也给一下吧

import numpy as np
import torch
import torch.nn as nn
import util_trt
import glob,os,cv2

BATCH_SIZE = 1
BATCH = 79
height = 640
width = 640
CALIB_IMG_DIR = '/content/drive/MyDrive/yolov5/ikunData/images'
onnx_model_path = "runs/train/exp4/weights/FP32.onnx"
def preprocess_v1(image_raw):
    h, w, c = image_raw.shape
    image = cv2.cvtColor(image_raw, cv2.COLOR_BGR2RGB)
    # Calculate widht and height and paddings
    r_w = width / w
    r_h = height / h
    if r_h > r_w:
        tw = width
        th = int(r_w * h)
        tx1 = tx2 = 0
        ty1 = int((height - th) / 2)
        ty2 = height - th - ty1
    else:
        tw = int(r_h * w)
        th = height
        tx1 = int((width - tw) / 2)
        tx2 = width - tw - tx1
        ty1 = ty2 = 0
    # Resize the image with long side while maintaining ratio
    image = cv2.resize(image, (tw, th))
    # Pad the short side with (128,128,128)
    image = cv2.copyMakeBorder(
        image, ty1, ty2, tx1, tx2, cv2.BORDER_CONSTANT, (128, 128, 128)
    )
    image = image.astype(np.float32)
    # Normalize to [0,1]
    image /= 255.0
    # HWC to CHW format:
    image = np.transpose(image, [2, 0, 1])
    # CHW to NCHW format
    #image = np.expand_dims(image, axis=0)
    # Convert the image to row-major order, also known as "C order":
    #image = np.ascontiguousarray(image)
    return image


def preprocess(img):
    img = cv2.resize(img, (640, 640))
    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_v1(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 = "runs/train/exp4/weights/int8.engine"
    calibration_table = 'yolov5_tensorrt_int8_tools/models_save/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)
    assert engine_fixed, 'Broken engine_fixed'
    print('*** onnx to tensorrt completed ***\n')
    
if __name__ == '__main__':
    main()
    

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.coloradmin.cn/o/1199714.html

如若内容造成侵权/违法违规/事实不符,请联系多彩编程网进行投诉反馈,一经查实,立即删除!

相关文章

Technology Strategy Patterns 学习笔记8- Communicating the Strategy-Decks(ppt模板)

1 Ghost Deck/Blank Deck 1.1 It’s a special way of making an initial deck that has a certain purpose 1.2 you’re making sure you have figured out what all the important shots are before incurring the major expense of shooting them 1.3 需要从技术、战略、产…

2023 年最新企业微信官方会话机器人开发详细教程(更新中)

目标是开发一个简易机器人&#xff0c;能接收消息并作出回复。 获取企业 ID 企业信息页面链接地址&#xff1a;https://work.weixin.qq.com/wework_admin/frame#profile 自建企业微信机器人 配置机器人应用详情 功能配置 接收消息服务器配置 配置消息服务器配置 配置环境变量…

[01]汇川IMC30G-E系列运动控制卡应用笔记

简介 IMC30G-E系列产品是汇川技术自主研制的高性能EtherCAT网络型运动控制器&#xff08;卡&#xff09;&#xff0c;同时兼容脉冲轴的控制&#xff1b;IMC30G-E支持点位/JOG、插补、多轴同步、高速位置比较输出、PWM等全面的运动控制功能&#xff0c;具备高同步控制精度。 开发…

OpenWRT浅尝 / 基于RAVPower-WD009便携路由文件宝的旁路网关配置

目录 前言需求分析手头的设备家庭网络拓扑图旁路网关配置OpenWRT固件选择OpenWRT固件刷入旁路网关配置流程 旁路网关的使用前置工作日常存储/关键备份内网穿透24小时待命下载器 前言 近期由于个人需求&#xff0c;需要一台OpenWRT设备实现一些功能。所以本文主要还是为了自己后…

k8s-实验部署 1

1、k8s集群部署 更改所有主机名称和解析 开启四台实验主机&#xff0c;k8s1 仓库&#xff1b;k8s2 集群控制节点&#xff1b; k8s3 和k8s4集群工作节点&#xff1b; 集群环境初始化 使用k8s1作为仓库&#xff0c;将所有的镜像都保存在本地&#xff0c;不要将集群从外部走 仓库…

金和OA jc6 任意文件上传漏洞复现

0x01 产品简介 金和OA协同办公管理系统软件&#xff08;简称金和OA&#xff09;&#xff0c;本着简单、适用、高效的原则&#xff0c;贴合企事业单位的实际需求&#xff0c;实行通用化、标准化、智能化、人性化的产品设计&#xff0c;充分体现企事业单位规范管理、提高办公效率…

学习率范围测试(LR Finder)脚本

简介 深度学习中的学习率是模型训练中至关重要的超参数之一。合适的学习率可以加速模型的收敛&#xff0c;提高训练效率&#xff0c;而不恰当的学习率可能导致训练过慢或者无法收敛。为了找到合适的学习率&#xff0c;LR Finder成为了一种强大的工具。 学习率范围测试&#x…

Django的ORM操作

文章目录 1.ORM操作1.1 表结构1.1.1 常见字段和参数1.1.2 表关系 2.ORM2.1 基本操作2.2 连接数据库2.3 基础增删改查2.3.1 增加2.3.2 查找2.3.4 删除2.3.4 修改 1.ORM操作 orm&#xff0c;关系对象映射&#xff0c;本质翻译的。 1.1 表结构 实现&#xff1a;创建表、修改表、…

思维模型 暗示效应

本系列文章 主要是 分享 思维模型&#xff0c;涉及各个领域&#xff0c;重在提升认知。无形中引导他人的思想和行为。 1 暗示效应的应用 1.1 暗示效应在商业品牌树立中的应用 可口可乐的品牌形象&#xff1a;可口可乐通过广告、包装和营销活动&#xff0c;向消费者传递了一种…

【递归】求根节点到叶节点数字之和(Java版)

目录 1.题目解析 2.讲解算法原理 3.代码 1.题目解析 LCR 049. 求根节点到叶节点数字之和 给定一个二叉树的根节点 root &#xff0c;树中每个节点都存放有一个 0 到 9 之间的数字。 每条从根节点到叶节点的路径都代表一个数字&#xff1a; 例如&#xff0c;从根节点到叶节点…

伙伴(buddy)系统原理

一、伙伴算法的由来 在实际情况中&#xff0c;操作系统必须能够在任意时刻申请和释放任意大小的内存&#xff0c;该函数的实现需要考虑延时问题和碎片问题。 延时问题指的是系统查找到可分配单元的时间变长&#xff0c;例如程序请求分配一个64KB的内存空间&#xff0c;系统查看…

Technology Strategy Patterns 学习笔记9 - bringing it all together

1 Patterns Map 2 Creating the Strategy 2.1 Ansoff Growth Matrix 和owth-share Matrix 区别参见https://fourweekmba.com/bcg-matrix-vs-ansoff-matrix/ 3 Communicating

STM32F407: CMSIS-DSP库的移植(基于库文件)

目录 1. 源码下载 2. DSP库源码简介 3.基于库的移植(DSP库的使用) 3.1 实验1 3.2 实验2 4. 使用V6版本的编译器进行编译 上一篇&#xff1a;STM32F407-Discovery的硬件FPU-CSDN博客 1. 源码下载 Github地址&#xff1a;GitHub - ARM-software/CMSIS_5: CMSIS Version 5…

开发者测试2023省赛--Square测试用例

测试结果 官方提交结果 EclEmma PITest 被测文件 [1/7] Square.java /*** This class implements the Square block cipher.** <P>* <b>References</b>** <P>* The Square algorithm was developed by <a href="mailto:Daemen.J@banksys.co…

AWS云服务器EC2实例进行操作系统迁移

AWS云服务器EC2实例进行操作系统迁移 文章目录 AWS云服务器EC2实例进行操作系统迁移1. 亚马逊EC2云服务器简介1.2 亚马逊EC2云务器与弹性云服务器区别 2. 亚马逊EC2云服务器配置流程2.1 亚马逊EC2云服务器实例配置2.1.1 EC2实例购买教程2.1.1 EC2实例初始化配置2.1.2 远程登录E…

Gold-YOLO:基于收集-分配机制的高效目标检测器

文章目录 摘要1、简介2、相关工作2.1、实时目标检测器2.2、基于Transformer的目标检测2.3、用于目标检测的多尺度特征 3、方法3.1、预备知识3.2、低级收集和分发分支3.3、高阶段收集和分发分支3.4、增强的跨层信息流3.5、遮罩图像建模预训练 4、实验4.1、设置4.2、比较4.3.2、 …

recycleView(三)动态修改背景色

效果图 1.关键代码 1. // 定义一个变量来记录滑动的距离var scrollDistance 0// 在RecycleView的滑动监听器中更新滑动的距离binding.recyclerView.addOnScrollListener(object : RecyclerView.OnScrollListener() {override fun onScrolled(recyclerView: RecyclerView, …

Pinia 状态管理器 菠萝

Pinia介绍&#xff1a; Pinia 是 Vue 的专属状态管理库&#xff0c;它允许你跨组件或页面共享状态。 Pinia 大小只有 1kb 左右&#xff0c;超轻量级&#xff0c;你甚至可能忘记它的存在&#xff01; 相比 Vuex,Pinia 的优点&#xff1a; 更贴合 Vue 3 的 Composition API 风…

Leetcode—191.位1的个数【简单】

2023每日刷题&#xff08;二十七&#xff09; Leetcode—191.位1的个数 实现代码 int hammingWeight(uint32_t n) {int ans 0;for(int i 0; i < 32; i) {if(n & ((long long)1 << i)) {ans;}}return ans; }运行结果 翻转比特1思路 就解法一的代码实现来说&am…

基于Transformer架构的ChatGPT:三步带你了解它的工作原理

作者&#xff1a;Insist-- 个人主页&#xff1a;insist--个人主页 梦想从未散场&#xff0c;传奇永不落幕&#xff0c;博主会持续更新优质网络知识、Python知识、Linux知识以及各种小技巧&#xff0c;愿你我共同在CSDN进步 目录 一、Transformer架构 1. 自注意力层 2. 前馈神…