├── README.md ├── __pycache__ ├── calibrator.cpython-36.pyc ├── util_trt.cpython-36.pyc └── util_trt.cpython-38.pyc ├── calibrator.py ├── convert_trt_quant.py ├── models_save ├── nanodet_calibration.cache ├── nanodet_int8.trt ├── yolov5s_calibration.cache └── yolov5s_int8.trt └── util_trt.py /README.md: -------------------------------------------------------------------------------- 1 | # 环境配置 2 | 3 | ubuntu：18.04 4 | 5 | cuda：11.0 6 | 7 | cudnn：8.0 8 | 9 | tensorrt：7.2.16 10 | 11 | OpenCV：3.4.2 12 | 13 | cuda，cudnn，tensorrt和OpenCV安装包（编译好了，也可以自己从官网下载编译）可以从链接: https://pan.baidu.com/s/1dpMRyzLivnBAca2c_DIgGw 密码: 0rct 14 | 15 | cuda安装 16 | 17 | 如果系统有安装驱动，运行如下命令卸载 18 | 19 | sudo apt-get purge nvidia* 20 | 21 | 禁用nouveau，运行如下命令 22 | 23 | sudo vim /etc/modprobe.d/blacklist.conf 24 | 25 | 在末尾添加 26 | 27 | blacklist nouveau 28 | 29 | 然后执行 30 | 31 | sudo update-initramfs -u 32 | 33 | chmod +x cuda_11.0.2_450.51.05_linux.run 34 | 35 | sudo ./cuda_11.0.2_450.51.05_linux.run 36 | 37 | 是否接受协议: accept 38 | 39 | 然后选择Install 40 | 41 | 最后回车 42 | 43 | vim ~/.bashrc 添加如下内容： 44 | 45 | export PATH=/usr/local/cuda-11.0/bin:$PATH 46 | 47 | export LD_LIBRARY_PATH=/usr/local/cuda-11.0/lib64:$LD_LIBRARY_PATH 48 | 49 | source .bashrc 激活环境 50 | 51 | cudnn 安装 52 | 53 | tar -xzvf cudnn-11.0-linux-x64-v8.0.4.30.tgz 54 | 55 | cd cuda/include 56 | 57 | sudo cp *.h /usr/local/cuda-11.0/include 58 | 59 | cd cuda/lib64 60 | 61 | sudo cp libcudnn* /usr/local/cuda-11.0/lib64 62 | 63 | tensorrt及OpenCV安装 64 | 65 | 定位到用户根目录 66 | 67 | tar -xzvf TensorRT-7.2.1.6.Ubuntu-18.04.x86_64-gnu.cuda-11.0.cudnn8.0.tar.gz 68 | 69 | cd TensorRT-7.2.1.6/python，该目录有4个python版本的tensorrt安装包 70 | 71 | sudo pip3 install tensorrt-7.2.1.6-cp37-none-linux_x86_64.whl（根据自己的python版本安装） 72 | 73 | pip install pycuda 安装python版本的cuda 74 | 75 | 定位到用户根目录 76 | 77 | tar -xzvf opencv-3.4.2.zip 以备推理调用 78 | 79 | # yolov5s模型转换onnx 80 | 81 | pip install onnx 82 | 83 | pip install onnx-simplifier 84 | 85 | git clone https://github.com/ultralytics/yolov5.git 86 | 87 | cd yolov5/models 88 | 89 | vim common.py 90 | 91 | 把BottleneckCSP类下的激活函数替换为relu，tensorrt对leakyRelu int8量化不稳定（这是一个深坑，大家记得避开）即修改为self.act = nn.ReLU(inplace=True) 92 | 93 | 训练得到模型后 94 | 95 | cd yolov5 96 | 97 | python models/export.py --weights 训练得到的模型权重路径 --img-size 训练图片输入尺寸 98 | 99 | python3 -m onnxsim onnx模型名称 yolov5s-simple.onnx 得到最终简化后的onnx模型 100 | 101 | # onnx模型转换为 int8 tensorrt引擎 102 | 103 | git clone https://github.com/Wulingtian/yolov5_tensorrt_int8_tools.git（求star） 104 | 105 | cd yolov5_tensorrt_int8_tools 106 | 107 | vim convert_trt_quant.py 修改如下参数 108 | 109 | BATCH_SIZE 模型量化一次输入多少张图片 110 | 111 | BATCH 模型量化次数 112 | 113 | height width 输入图片宽和高 114 | 115 | CALIB_IMG_DIR 训练图片路径，用于量化 116 | 117 | onnx_model_path onnx模型路径 118 | 119 | python convert_trt_quant.py 量化后的模型存到models_save目录下 120 | 121 | # tensorrt模型推理 122 | 123 | git clone https://github.com/Wulingtian/yolov5_tensorrt_int8.git（求star） 124 | 125 | cd yolov5_tensorrt_int8 126 | 127 | vim CMakeLists.txt 128 | 129 | 修改USER_DIR参数为自己的用户根目录 130 | 131 | vim http://yolov5s_infer.cc 修改如下参数 132 | 133 | output_name1 output_name2 output_name3 yolov5模型有3个输出 134 | 135 | 我们可以通过netron查看模型输出名 136 | 137 | pip install netron 安装netron 138 | 139 | vim netron_yolov5s.py 把如下内容粘贴 140 | 141 | import netron 142 | 143 | netron.start('此处填充简化后的onnx模型路径', port=3344) 144 | 145 | python netron_yolov5s.py 即可查看 模型输出名 146 | 147 | trt_model_path 量化的的tensorrt推理引擎（models_save目录下trt后缀的文件） 148 | 149 | test_img 测试图片路径 150 | 151 | INPUT_W INPUT_H 输入图片宽高 152 | 153 | NUM_CLASS 训练的模型有多少类 154 | 155 | NMS_THRESH nms阈值 156 | 157 | CONF_THRESH 置信度 158 | 159 | 参数配置完毕 160 | 161 | mkdir build 162 | 163 | cd build 164 | 165 | cmake .. 166 | 167 | make 168 | 169 | ./YoloV5sEngine 输出平均推理时间，以及保存预测图片到当前目录下，至此，部署完成！ 170 | -------------------------------------------------------------------------------- /__pycache__/calibrator.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Wulingtian/yolov5_tensorrt_int8_tools/cfbca437adc942903973838810a840faacd0b7d7/__pycache__/calibrator.cpython-36.pyc -------------------------------------------------------------------------------- /__pycache__/util_trt.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Wulingtian/yolov5_tensorrt_int8_tools/cfbca437adc942903973838810a840faacd0b7d7/__pycache__/util_trt.cpython-36.pyc -------------------------------------------------------------------------------- /__pycache__/util_trt.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Wulingtian/yolov5_tensorrt_int8_tools/cfbca437adc942903973838810a840faacd0b7d7/__pycache__/util_trt.cpython-38.pyc -------------------------------------------------------------------------------- /calibrator.py: -------------------------------------------------------------------------------- 1 | # *** tensorrt校准模块 *** 2 | 3 | import os 4 | import torch 5 | import torch.nn.functional as F 6 | import tensorrt as trt 7 | import pycuda.driver as cuda 8 | import pycuda.autoinit 9 | import numpy as np 10 | import ctypes 11 | import logging 12 | import util_trt 13 | logger = logging.getLogger(__name__) 14 | ctypes.pythonapi.PyCapsule_GetPointer.restype = ctypes.c_char_p 15 | ctypes.pythonapi.PyCapsule_GetPointer.argtypes = [ctypes.py_object, ctypes.c_char_p] 16 | 17 | # calibrator 18 | #IInt8EntropyCalibrator2 19 | #IInt8LegacyCalibrator 20 | #IInt8EntropyCalibrator 21 | #IInt8MinMaxCalibrator 22 | class Calibrator(trt.IInt8EntropyCalibrator): 23 | def __init__(self, stream, cache_file=""): 24 | trt.IInt8EntropyCalibrator.__init__(self) 25 | self.stream = stream 26 | self.d_input = cuda.mem_alloc(self.stream.calibration_data.nbytes) 27 | self.cache_file = cache_file 28 | stream.reset() 29 | 30 | def get_batch_size(self): 31 | return self.stream.batch_size 32 | 33 | def get_batch(self, names): 34 | #print("############################################################") 35 | #print(names) 36 | #print("############################################################") 37 | batch = self.stream.next_batch() 38 | if not batch.size: 39 | return None 40 | 41 | cuda.memcpy_htod(self.d_input, batch) 42 | 43 | return [int(self.d_input)] 44 | 45 | def read_calibration_cache(self): 46 | # If there is a cache, use it instead of calibrating again. Otherwise, implicitly return None. 47 | if os.path.exists(self.cache_file): 48 | with open(self.cache_file, "rb") as f: 49 | logger.info("Using calibration cache to save time: {:}".format(self.cache_file)) 50 | return f.read() 51 | 52 | def write_calibration_cache(self, cache): 53 | with open(self.cache_file, "wb") as f: 54 | logger.info("Caching calibration data for future use: {:}".format(self.cache_file)) 55 | f.write(cache) 56 | -------------------------------------------------------------------------------- /convert_trt_quant.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import util_trt 5 | import glob,os,cv2 6 | 7 | BATCH_SIZE = 16 8 | BATCH = 100 9 | height = 640 10 | width = 640 11 | CALIB_IMG_DIR = '/home/willer/yolov5-3.1/data/coco/images/train2017' 12 | onnx_model_path = "/home/willer/yolov5-4.0/models/models_silu/yolov5s-simple.onnx" 13 | def preprocess_v1(image_raw): 14 | h, w, c = image_raw.shape 15 | image = cv2.cvtColor(image_raw, cv2.COLOR_BGR2RGB) 16 | # Calculate widht and height and paddings 17 | r_w = width / w 18 | r_h = height / h 19 | if r_h > r_w: 20 | tw = width 21 | th = int(r_w * h) 22 | tx1 = tx2 = 0 23 | ty1 = int((height - th) / 2) 24 | ty2 = height - th - ty1 25 | else: 26 | tw = int(r_h * w) 27 | th = height 28 | tx1 = int((width - tw) / 2) 29 | tx2 = width - tw - tx1 30 | ty1 = ty2 = 0 31 | # Resize the image with long side while maintaining ratio 32 | image = cv2.resize(image, (tw, th)) 33 | # Pad the short side with (128,128,128) 34 | image = cv2.copyMakeBorder( 35 | image, ty1, ty2, tx1, tx2, cv2.BORDER_CONSTANT, (128, 128, 128) 36 | ) 37 | image = image.astype(np.float32) 38 | # Normalize to [0,1] 39 | image /= 255.0 40 | # HWC to CHW format: 41 | image = np.transpose(image, [2, 0, 1]) 42 | # CHW to NCHW format 43 | #image = np.expand_dims(image, axis=0) 44 | # Convert the image to row-major order, also known as "C order": 45 | #image = np.ascontiguousarray(image) 46 | return image 47 | 48 | 49 | def preprocess(img): 50 | img = cv2.resize(img, (640, 640)) 51 | img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 52 | img = img.transpose((2, 0, 1)).astype(np.float32) 53 | img /= 255.0 54 | return img 55 | 56 | class DataLoader: 57 | def __init__(self): 58 | self.index = 0 59 | self.length = BATCH 60 | self.batch_size = BATCH_SIZE 61 | # self.img_list = [i.strip() for i in open('calib.txt').readlines()] 62 | self.img_list = glob.glob(os.path.join(CALIB_IMG_DIR, "*.jpg")) 63 | 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' 64 | print('found all {} images to calib.'.format(len(self.img_list))) 65 | self.calibration_data = np.zeros((self.batch_size,3,height,width), dtype=np.float32) 66 | 67 | def reset(self): 68 | self.index = 0 69 | 70 | def next_batch(self): 71 | if self.index < self.length: 72 | for i in range(self.batch_size): 73 | assert os.path.exists(self.img_list[i + self.index * self.batch_size]), 'not found!!' 74 | img = cv2.imread(self.img_list[i + self.index * self.batch_size]) 75 | img = preprocess_v1(img) 76 | self.calibration_data[i] = img 77 | 78 | self.index += 1 79 | 80 | # example only 81 | return np.ascontiguousarray(self.calibration_data, dtype=np.float32) 82 | else: 83 | return np.array([]) 84 | 85 | def __len__(self): 86 | return self.length 87 | 88 | def main(): 89 | # onnx2trt 90 | fp16_mode = False 91 | int8_mode = True 92 | print('*** onnx to tensorrt begin ***') 93 | # calibration 94 | calibration_stream = DataLoader() 95 | engine_model_path = "models_save/yolov5s_int8.trt" 96 | calibration_table = 'models_save/yolov5s_calibration.cache' 97 | # fixed_engine,校准产生校准表 98 | engine_fixed = util_trt.get_engine(BATCH_SIZE, onnx_model_path, engine_model_path, fp16_mode=fp16_mode, 99 | int8_mode=int8_mode, calibration_stream=calibration_stream, calibration_table_path=calibration_table, save_engine=True) 100 | assert engine_fixed, 'Broken engine_fixed' 101 | print('*** onnx to tensorrt completed ***\n') 102 | 103 | if __name__ == '__main__': 104 | main() 105 | 106 | -------------------------------------------------------------------------------- /models_save/nanodet_calibration.cache: -------------------------------------------------------------------------------- 1 | TRT-7201-EntropyCalibration 2 | input.1: 3c010a14 3 | 804: 3c2bc75d 4 | 432: 3c16645c 5 | 433: 3bb66d00 6 | 807: 3b34643a 7 | 810: 3b6dcc49 8 | 438: 3b7b5d57 9 | 813: 3b8c98e9 10 | 441: 3b66518f 11 | 816: 3b5b7d09 12 | 819: 3bc02823 13 | 446: 3bc953a0 14 | 450: 3bc9e85f 15 | 452: 3bc9e85f 16 | 453: 3b883866 17 | 454: 3bc9e85f 18 | 822: 3b6539b9 19 | 457: 3b64a61c 20 | 825: 3b942d9c 21 | 828: 3b85a769 22 | 462: 3b876365 23 | 466: 3b883866 24 | 468: 3b883866 25 | 469: 3b341eba 26 | 470: 3b883866 27 | 831: 3b6dc0cd 28 | 473: 3b667e55 29 | 834: 3bf4e390 30 | 837: 3bf46014 31 | 478: 3bbc4517 32 | 482: 3b9c3955 33 | 484: 3b9c3955 34 | 485: 3b9ed50c 35 | 486: 3b7b8f84 36 | 840: 3b7c074c 37 | 489: 3b762981 38 | 843: 3be9fd30 39 | 846: 3c1a20d8 40 | 494: 3bf1c5b6 41 | 498: 3bafcc5c 42 | 500: 3bafcc5c 43 | 849: 3b9d7b9c 44 | 852: 3b9b8f74 45 | 505: 3b8b651b 46 | 855: 3b844b2f 47 | 508: 3b2d5a06 48 | 858: 3c249891 49 | 861: 3bcfa344 50 | 513: 3bb7a491 51 | 517: 3bbb75d1 52 | 519: 3bbb75d1 53 | 520: 3ba48e55 54 | 521: 3b8926f0 55 | 864: 3b733403 56 | 524: 3b690b3c 57 | 867: 3b9c2946 58 | 870: 3b8f540f 59 | 529: 3b8f540f 60 | 533: 3bbf470f 61 | 535: 3bbf470f 62 | 536: 3b92d5b8 63 | 537: 3b8f03de 64 | 873: 3b66b18e 65 | 540: 3b34e507 66 | 876: 3bd4c876 67 | 879: 3be7110d 68 | 545: 3b9c6f7c 69 | 549: 3b9a408d 70 | 551: 3b9a408d 71 | 552: 3bab4838 72 | 553: 3ba4718c 73 | 882: 3b8ff697 74 | 556: 3b476516 75 | 885: 3c08ba6f 76 | 888: 3bfc622b 77 | 561: 3bb7c756 78 | 565: 3bba20f6 79 | 567: 3bba20f6 80 | 568: 3bbf8447 81 | 569: 3baf364d 82 | 891: 3b9448a4 83 | 572: 3b4bb139 84 | 894: 3c0ded6c 85 | 897: 3c046d63 86 | 577: 3bc4665d 87 | 581: 3bc560b6 88 | 583: 3bc560b6 89 | 584: 3b966ff9 90 | 585: 3bbdfe67 91 | 900: 3b895194 92 | 588: 3b498225 93 | 903: 3bff9304 94 | 906: 3bfb8787 95 | 593: 3b87d95a 96 | 597: 3b8f9788 97 | 599: 3b8f9788 98 | 600: 3b939423 99 | 601: 3b947b47 100 | 909: 3ba124b5 101 | 604: 3b55aa3c 102 | 912: 3c0ceb8a 103 | 915: 3c17e242 104 | 609: 3ba9c907 105 | 613: 3ba9c907 106 | 615: 3ba9c907 107 | 616: 3b9a30ed 108 | 617: 3ba5d8c2 109 | 918: 3b97f0e4 110 | 620: 3b626b69 111 | 921: 3c1244ea 112 | 924: 3c2884c6 113 | 625: 3b94c13f 114 | 629: 3b9e8262 115 | 631: 3b9e8262 116 | 927: 3bbf7cd3 117 | 930: 3bca5c5b 118 | 636: 3b814130 119 | 933: 3baa92e1 120 | 639: 3b0baea4 121 | 936: 3c029f6a 122 | 939: 3bdeee08 123 | 644: 3b4b4002 124 | 648: 3b7d2601 125 | 650: 3b7d2601 126 | 651: 3b529bf2 127 | 652: 3b7e023f 128 | 942: 3b82c616 129 | 655: 3b06a8ed 130 | 945: 3bef047d 131 | 948: 3bcab145 132 | 660: 3b665071 133 | 664: 3b51028f 134 | 666: 3b51028f 135 | 667: 3b4e21aa 136 | 668: 3b65ea44 137 | 951: 3b72060c 138 | 671: 3b2d1011 139 | 954: 3be8d30d 140 | 957: 3bd7a823 141 | 676: 3b5babaa 142 | 680: 3b585dba 143 | 682: 3b585dba 144 | 683: 3b515470 145 | 684: 3b5b88bc 146 | 960: 3b81aa47 147 | 687: 3b1329f7 148 | 963: 3c0551e2 149 | 966: 3c01b10f 150 | 692: 3b7cf2da 151 | 696: 3b6be038 152 | 698: 3b6be038 153 | 701: 3b6939bf 154 | 702: 3aee4eaa 155 | 703: 3bc2403c 156 | 704: 3b5b726d 157 | 705: 3bcd4f70 158 | 707: 3c0ac85a 159 | 709: 3bb55cb1 160 | 969: 3c50f724 161 | 712: 3c03ce32 162 | 972: 3bae0e44 163 | 715: 3bd2b7a7 164 | 975: 3c4cec4b 165 | 718: 3c2a232b 166 | 978: 3d0b2f43 167 | 721: 3cbcf8f2 168 | 722: 3d710269 169 | 723: 3d5b9ee4 170 | 724: 3d710269 171 | 725: 3ad7c149 172 | 728: 3ad7c149 173 | 731: 3d710269 174 | (Unnamed Layer* 210) [Shuffle]_output: 3d710269 175 | (Unnamed Layer* 211) [Softmax]_output: 3b9a51ad 176 | 734: 3b9a51ad 177 | (Unnamed Layer* 214) [Shuffle]_output: 3d61d1a3 178 | 736: 3d225fb2 179 | 737: 3d225fb2 180 | 739: 3d225fb2 181 | 981: 3c18b6d7 182 | 742: 3bcda098 183 | 984: 3bcd9eb1 184 | 745: 3b3c0ddb 185 | 987: 3c0d0c38 186 | 748: 3b9fb31e 187 | 990: 3d2961b1 188 | 751: 3cb9c1ff 189 | 752: 3d59a0e6 190 | 753: 3d5a6908 191 | 754: 3d59951f 192 | 755: 3a136e3d 193 | 758: 3a136e3d 194 | 761: 3d59951f 195 | (Unnamed Layer* 237) [Shuffle]_output: 3d59951f 196 | (Unnamed Layer* 238) [Softmax]_output: 3b9ed660 197 | 764: 3b9ed660 198 | (Unnamed Layer* 241) [Shuffle]_output: 3d61d1a3 199 | 766: 3d028d2d 200 | 767: 3d028d2d 201 | 769: 3d028d2d 202 | 993: 3bc1fb27 203 | 772: 3ba9a3ea 204 | 996: 3b713346 205 | 775: 3b0df498 206 | 999: 3bd36c00 207 | 778: 3b5b790e 208 | 1002: 3d5b231f 209 | 781: 3cb49715 210 | 782: 3d2819ff 211 | 783: 3d2c310e 212 | 784: 3d26b0a4 213 | 785: 3a0f4e88 214 | 788: 3a0f4e88 215 | 791: 3d26b0a4 216 | (Unnamed Layer* 264) [Shuffle]_output: 3d26b0a4 217 | (Unnamed Layer* 265) [Softmax]_output: 3b051d63 218 | 794: 3b051d63 219 | (Unnamed Layer* 268) [Shuffle]_output: 3d61d1a3 220 | 796: 3cdfe9df 221 | 797: 3cdfe9df 222 | 799: 3cdfe9df 223 | -------------------------------------------------------------------------------- /models_save/nanodet_int8.trt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Wulingtian/yolov5_tensorrt_int8_tools/cfbca437adc942903973838810a840faacd0b7d7/models_save/nanodet_int8.trt -------------------------------------------------------------------------------- /models_save/yolov5s_calibration.cache: -------------------------------------------------------------------------------- 1 | TRT-7201-EntropyCalibration 2 | images: 3c010a14 3 | 127: 3c7fb6aa 4 | 129: 3c79ec7a 5 | 130: 3e109797 6 | 132: 3dd57f1d 7 | 133: 3d4bb166 8 | 135: 3c1f188e 9 | 136: 3d14e375 10 | 138: 3cb7d4ca 11 | 139: 3d158eaf 12 | 142: 3d53c94e 13 | 143: 3d628152 14 | 145: 3d1f1415 15 | 147: 3d4707fb 16 | 149: 3c5fa90e 17 | 150: 3d209f08 18 | 152: 3ca500fb 19 | 153: 3c340555 20 | 155: 3c0c4099 21 | 156: 3d122ce3 22 | 158: 3d0e544a 23 | 159: 3cdd8ae9 24 | 162: 3c8a9820 25 | 163: 3cf3085f 26 | 165: 3c802d91 27 | 166: 3cd04b32 28 | 169: 3cb06f4f 29 | 170: 3d068bdc 30 | 172: 3cd2925b 31 | 173: 3d25c831 32 | 176: 3d25e3c3 33 | 177: 3d28e748 34 | 179: 3c622f3e 35 | 181: 3d5193c1 36 | 183: 3c80c0a5 37 | 184: 3d2db50d 38 | 186: 3cb31e04 39 | 187: 3cc77fef 40 | 189: 3c084596 41 | 190: 3d37af47 42 | 192: 3cb23edd 43 | 193: 3d16096d 44 | 196: 3ca3d3d4 45 | 197: 3d459c32 46 | 199: 3cc021f7 47 | 200: 3d310cf5 48 | 203: 3d29d6d7 49 | 204: 3d56bfc5 50 | 206: 3c8fdd33 51 | 207: 3d9fa8b2 52 | 210: 3d82da12 53 | 211: 3d2eabf1 54 | 213: 3d0052cd 55 | 215: 3d60773f 56 | 217: 3c8863c2 57 | 218: 3d824921 58 | 220: 3cd3bda6 59 | 221: 3d4a4e76 60 | 223: 3d27ad3d 61 | 224: 3d63a75f 62 | 225: 3d62754f 63 | 226: 3d85c0b9 64 | 228: 3d748ce2 65 | 230: 3c914f59 66 | 231: 3d87504b 67 | 233: 3cea1a9f 68 | 234: 3d82a532 69 | 236: 3d0160ad 70 | 237: 3d899a78 71 | 239: 3d008d92 72 | 240: 3d684027 73 | 242: 3c8fae36 74 | 244: 3da4c706 75 | 246: 3d132d2a 76 | 247: 3ddc0e0d 77 | 249: 3d121dac 78 | 254: 3d121dac 79 | 256: 3d89f8b5 80 | 258: 3cb4bc8f 81 | 259: 3d99c69c 82 | 261: 3cc242e4 83 | 262: 3d849480 84 | 264: 3c903159 85 | 265: 3dad7485 86 | 267: 3cd19a0c 87 | 269: 3dcdacf2 88 | 271: 3d2762b3 89 | 272: 3db9ec9d 90 | 274: 3d939d64 91 | 279: 3d939d64 92 | 281: 3d792d77 93 | 283: 3c51bc3e 94 | 284: 3d79b464 95 | 286: 3ce6fb60 96 | 287: 3dabf5ba 97 | 289: 3d43c5fa 98 | 290: 3db89991 99 | 292: 3dbde1ee 100 | 294: 3e94e3d9 101 | 296: 3e17bb67 102 | 297: 3ddde51e 103 | 299: 3d927b29 104 | 301: 3dbd7038 105 | 303: 3d263370 106 | 304: 3d855c72 107 | 306: 3cd2c9ee 108 | 307: 3df7053a 109 | 309: 3d940b3d 110 | 310: 3de2fb0f 111 | 312: 3e0981cd 112 | 314: 3e7467dd 113 | 316: 3dc66f36 114 | 317: 3dda066e 115 | 319: 3d687ee3 116 | 321: 3dc95dae 117 | 323: 3d3fa7c1 118 | 324: 3deadc04 119 | 326: 3d9422af 120 | 327: 3e134be9 121 | 329: 3dc0ec32 122 | 330: 3dcbb8e7 123 | 332: 3dbfaee9 124 | 334: 3e130d3f 125 | 336: 3d91ca8c 126 | 337: 3e7d5e91 127 | output: 3e7d5e91 128 | 357: 3eb2f27a 129 | 376: 3eb2f27a 130 | 377: 3e8ccce7 131 | 396: 3e8ccce7 132 | -------------------------------------------------------------------------------- /models_save/yolov5s_int8.trt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Wulingtian/yolov5_tensorrt_int8_tools/cfbca437adc942903973838810a840faacd0b7d7/models_save/yolov5s_int8.trt -------------------------------------------------------------------------------- /util_trt.py: -------------------------------------------------------------------------------- 1 | # tensorrt-lib 2 | 3 | import os 4 | import tensorrt as trt 5 | import pycuda.autoinit 6 | import pycuda.driver as cuda 7 | from calibrator import Calibrator 8 | from torch.autograd import Variable 9 | import torch 10 | import numpy as np 11 | import time 12 | # add verbose 13 | TRT_LOGGER = trt.Logger(trt.Logger.VERBOSE) # ** engine可视化 ** 14 | 15 | # create tensorrt-engine 16 | # fixed and dynamic 17 | def get_engine(max_batch_size=1, onnx_file_path="", engine_file_path="",\ 18 | fp16_mode=False, int8_mode=False, calibration_stream=None, calibration_table_path="", save_engine=False): 19 | """Attempts to load a serialized engine if available, otherwise builds a new TensorRT engine and saves it.""" 20 | def build_engine(max_batch_size, save_engine): 21 | """Takes an ONNX file and creates a TensorRT engine to run inference with""" 22 | with trt.Builder(TRT_LOGGER) as builder, \ 23 | builder.create_network(1) as network,\ 24 | trt.OnnxParser(network, TRT_LOGGER) as parser: 25 | 26 | # parse onnx model file 27 | if not os.path.exists(onnx_file_path): 28 | quit('ONNX file {} not found'.format(onnx_file_path)) 29 | print('Loading ONNX file from path {}...'.format(onnx_file_path)) 30 | with open(onnx_file_path, 'rb') as model: 31 | print('Beginning ONNX file parsing') 32 | parser.parse(model.read()) 33 | assert network.num_layers > 0, 'Failed to parse ONNX model. \ 34 | Please check if the ONNX model is compatible ' 35 | print('Completed parsing of ONNX file') 36 | print('Building an engine from file {}; this may take a while...'.format(onnx_file_path)) 37 | 38 | # build trt engine 39 | builder.max_batch_size = max_batch_size 40 | builder.max_workspace_size = 1 << 30 # 1GB 41 | builder.fp16_mode = fp16_mode 42 | if int8_mode: 43 | builder.int8_mode = int8_mode 44 | assert calibration_stream, 'Error: a calibration_stream should be provided for int8 mode' 45 | builder.int8_calibrator = Calibrator(calibration_stream, calibration_table_path) 46 | print('Int8 mode enabled') 47 | engine = builder.build_cuda_engine(network) 48 | if engine is None: 49 | print('Failed to create the engine') 50 | return None 51 | print("Completed creating the engine") 52 | if save_engine: 53 | with open(engine_file_path, "wb") as f: 54 | f.write(engine.serialize()) 55 | return engine 56 | 57 | if os.path.exists(engine_file_path): 58 | # If a serialized engine exists, load it instead of building a new one. 59 | print("Reading engine from file {}".format(engine_file_path)) 60 | with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime: 61 | return runtime.deserialize_cuda_engine(f.read()) 62 | else: 63 | return build_engine(max_batch_size, save_engine) 64 | --------------------------------------------------------------------------------