将YOLOv8模型从PyTorch的.pt格式转换为TensorRT的.engine格式

TensorRT是由NVIDIA开发的一款高级软件开发套件(SDK)，专为高速深度学习推理而设计。它非常适合目标检测等实时应用。该工具包可针对NVIDIA GPU优化深度学习模型，从而实现更快、更高效的运行。TensorRT模型经过TensorRT优化，包括层融合(layer fusion)、精度校准(precision calibration)(INT8和FP16)、动态张量内存管理和内核自动调整(kernel auto-tuning)等技术。将深度学习模型转换为TensorRT格式可充分发挥NVIDIA GPU的潜力。

TensorRT可兼容各种模型格式，包括TensorFlow、PyTorch和ONNX。

TensorRT模型的主要特点：

(1).Precision Calibration：TensorRT支持精度校准，允许根据特定的精度要求对模型进行微调(fine-tuned)。这包括对INT8和FP16等精度较低的格式的支持，这可以在保持可接受的精度水平的同时进一步提高推理速度。

(2).Layer Fusion：TensorRT优化过程包括层融合，即将神经网络的多个层组合成一个操作。这通过最小化内存访问和计算来减少计算开销并提高推理速度。

(3).Dynamic Tensor Memory Management：TensorRT可有效管理推理过程中的张量内存使用情况，从而减少内存开销并优化内存分配。这可提高GPU内存利用率。

(4).Automatic Kernel Tuning：TensorRT采用自动内核调整，为模型的每一层选择最优化的GPU内核。这种自适应方法可确保模型充分利用GPU的计算能力。

TensorRT中的部署选项：

(1).Deploying within TensorFlow：此方法将TensorRT集成到TensorFlow中，使优化的模型可以在TensorFlow环境中运行。对于混合了受支持层和不受支持的层(a mix of supported and unsupported layers)的模型，此方法非常有用，因为TF-TRT可以高效处理这些层。

(2).Standalone TensorRT Runtime API：提供精细控制，非常适合性能关键型应用程序。它更复杂，但允许自定义实现不受支持的运算符。

(3).NVIDIA Triton Inference Server：支持各种框架模型的选项。它特别适合云端或边缘端推理(cloud or edge inference)，提供并发模型(concurrent model)执行和模型分析等功能。

训练生成TensorRT支持的.engine格式模型：

训练代码如下所示：

import argparse
import colorama
from ultralytics import YOLO
import torch

def parse_args():
	parser = argparse.ArgumentParser(description="YOLOv8 train")
	parser.add_argument("--yaml", required=True, type=str, help="yaml file")
	parser.add_argument("--epochs", required=True, type=int, help="number of training")
	parser.add_argument("--task", required=True, type=str, choices=["detect", "segment"], help="specify what kind of task")

	args = parser.parse_args()
	return args

def train(task, yaml, epochs):
	if task == "detect":
		model = YOLO("yolov8n.pt") # load a pretrained model
	elif task == "segment":
		model = YOLO("yolov8n-seg.pt") # load a pretrained model
	else:
		print(colorama.Fore.RED + "Error: unsupported task:", task)
		raise

	results = model.train(data=yaml, epochs=epochs, imgsz=640) # train the model

	metrics = model.val() # It'll automatically evaluate the data you trained, no arguments needed, dataset and settings remembered

	# model.export(format="onnx") #, dynamic=True) # export the model, cannot specify dynamic=True, opencv does not support
	model.export(format="onnx", opset=12, simplify=True, dynamic=False, imgsz=640)
	model.export(format="torchscript") # libtorch
	model.export(format="engine", imgsz=640, dynamic=False, verbose=False, batch=1, workspace=2) # tensorrt fp32
	# model.export(format="engine", imgsz=640, dynamic=True, verbose=True, batch=4, workspace=2, half=True) # tensorrt fp16
	# model.export(format="engine", imgsz=640, dynamic=True, verbose=True, batch=4, workspace=2, int8=True, data=yaml) # tensorrt int8

if __name__ == "__main__":
	# python test_yolov8_train.py --yaml datasets/melon_new_detect/melon_new_detect.yaml --epochs 1000 --task detect
	colorama.init()
	args = parse_args()

	if torch.cuda.is_available():
		print("Runging on GPU")
	else:
		print("Runting on CPU")

	train(args.task, args.yaml, args.epochs)

	print(colorama.Fore.GREEN + "====== execution completed ======")

使用INT8量化导出TensorRT：会执行训练后量化(post-training quantization, PTQ)，即在模型训练完成后，无需重新训练即可对模型进行量化。TensorRT 使用校准进行PTQ。

注：确保使用TensorRT模型权重进行部署的同一设备以INT8精度进行导出，因为校准结果可能因设备而异。

配置INT8导出：使用导出Ultralytics YOLO模型时提供的参数将极大地影响导出模型的性能。还需要根据可用的设备资源来选择它们，但是默认参数应该适用于大多数 Ampere(或更新版本)架构的NVIDIA独立GPU。使用的校准算法是"ENTROPY_CALIBRATION_2"。

workspace：控制转换模型权重时设备内存分配的大小(以GiB为单位)。

(1).根据校准需求和资源可用性调整workspace。虽然较大的workspace可能会增加校准时间，但它允许TensorRT探索更广泛的优化策略，从而有可能提高模型性能和准确性。相反，较小的workspace可以减少校准时间，但可能会限制优化策略，影响量化模型的质量。

(2).默认值workspace=4(GiB)，如果校准崩溃(没有警告就退出)，则可能需要增加此值。

(3).如果workspace的值大于设备可用的内存，TensorRT将在导出期间报告UNSUPPORTED_STATE，这意味着应该降低workspace的值。

(4).如果workspace设置为最大值并且校准失败/崩溃，请考虑减少imgsz和batch的值以减少内存要求。

切记：INT8的校准是针对每个设备的，借用"高端"GPU进行校准可能会导致在另一台设备上运行推理时性能不佳。

batch：用于推理的最大批次大小(batch-size)。推理期间可以使用较小的批次，但推理不会接受大于指定值的批次。

在校准过程中，将使用提供的两倍批次大小。使用小批次可能会导致校准过程中的缩放不准确。这是因为该过程会根据它看到的数据进行调整。小批次可能无法捕获整个值范围，从而导致最终校准出现问题，因此批次大小会自动加倍。如果没有指定批次大小batch=1，则校准将以batch=1*2 运行，以减少校准缩放错误。

NVIDIA的实验使他们建议使用至少500张代表模型数据的校准图像，并使用INT8量化校准。这是一个指导原则，而不是硬性要求，你需要试验哪些内容才能使你的数据集表现良好。由于使用TensorRT进行INT8校准需要校准数据，因此确保在TensorRT的int8=True时使用数据参数并使用data="my_dataset.yaml"，这将使用验证中的图像进行校准。当使用INT8量化导出到TensorRT时没有传递任何数据值时，默认将使用基于模型任务的"small"示例数据集之一，而不是抛出错误。

注：TensorRT将生成一个校准.cache，可以重复使用以加速使用相同数据导出未来模型权重，但当数据差异很大或批次值发生剧烈变化时，这可能会导致校准效果不佳。在这种情况下，应重命名现有.cache并将其移动到其他目录或完全删除。

将YOLO与TensorRT INT8结合使用的优势：

(1).减少模型大小：从FP32到INT8的量化可以将模型大小减小4倍(在磁盘或内存中)，从而缩短下载时间、降低存储要求并减少部署模型时的内存占用。

(2).更低功耗：INT8导出的YOLO模型的精度运算减少，与FP32模型相比，功耗更低，尤其是对于电池供电(battery-powered)的设备。

(3).提高推理速度：TensorRT针对目标硬件优化模型，可能提高GPU、嵌入式设备和加速器上的推理速度。

注：使用导出到TensorRT INT8的模型进行前几次推理调用时，预处理、推理和/或后处理时间(preprocessing, inference, and/or postprocessing times)可能会比平时更长。在推理过程中更改imgsz时也可能会出现这种情况，尤其是当imgsz与导出期间指定的值不同时(导出imgsz设置为TensorRT"最佳"配置文件)。

使用YOLO和TensorRT INT8的缺点：

(1).评估指标下降：使用较低的精度意味着mAP、精度、召回率或用于评估模型性能的任何其他指标可能会有所下降。

(2).增加开发时间：找到数据集和设备的INT8校准的"最佳"设置可能需要大量测试。

(3).硬件依赖性：校准和性能提升可能高度依赖于硬件，并且模型权重的可转移性较差。

TensorRT的性能改进可能因所使用的硬件而异。

注：以上文字描述主要来自：https://docs.ultralytics.com/integrations/tensorrt/

Windows10 Anaconda上配置TensorRT环境：

(1).配置Ultralytics CUDA开发环境，执行以下命令：

# install cuda 11.8
# install cudnn v8.7.0: copy the contents of bin,include,lib/x64 cudnn directories to the corresponding CUDA directories
conda create --name ultralytics-env-cuda python=3.8 -y
conda activate ultralytics-env-cuda
conda install -c pytorch -c nvidia -c conda-forge pytorch==2.2.2 torchvision==0.17.2 torchaudio==2.2.2 pytorch-cuda=11.8 ultralytics # pytorch 2.2.2

git clone https://github.com/fengbingchun/NN_Test
cd NN_Test/demo/Python

(2).从https://developer.nvidia.com/nvidia-tensorrt-8x-download 下载TensorRT 8.5 GA版本：TensorRT-8.5.3.1.Windows10.x86_64.cuda-11.8.cudnn8.6.zip，解压缩:

A.将bin、include目录下内容拷贝到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8对应目录下

B.将lib下的所有静态库拷贝到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8\lib\x64目录下

C.将lib下的所有动态库拷贝到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8\bin目录下

(3).进入到python目录，执行以下命令：

pip install tensorrt-8.5.3.1-cp38-none-win_amd64.whl

注：不能使用10.2 GA版本，否则会报Error: Unsupported SM: 0x601，在 https://docs.nvidia.com/deeplearning/tensorrt/release-notes/ 中有描述：NVIDIA Pascal (SM 6.x) devices are deprecated in TensorRT 8.6

注：无论指定是FP32、FP16还是INT8训练完生成的最终文件名都为best.engine，这里手动调整文件名

在网上下载了200多幅包含西瓜和冬瓜的图像组成melon数据集，使用生成的best.engine进行预测，代码如下所示：

import colorama
import argparse
from ultralytics import YOLO
import os
import torch

import numpy as np
np.bool = np.bool_ # Fix Error: AttributeError: module 'numpy' has no attribute 'bool'. OR: downgrade numpy: pip unistall numpy; pip install numpy==1.23.1

def parse_args():
	parser = argparse.ArgumentParser(description="YOLOv8 predict")
	parser.add_argument("--model", required=True, type=str, help="model file")
	parser.add_argument("--dir_images", required=True, type=str, help="directory of test images")
	parser.add_argument("--dir_result", required=True, type=str, help="directory where the image results are saved")

	args = parser.parse_args()
	return args

def get_images(dir):
	# supported image formats
	img_formats = (".bmp", ".jpeg", ".jpg", ".png", ".webp")
	images = []

	for file in os.listdir(dir):
		if os.path.isfile(os.path.join(dir, file)):
			# print(file)
			_, extension = os.path.splitext(file)
			for format in img_formats:
				if format == extension.lower():
					images.append(file)
					break

	return images

def predict(model, dir_images, dir_result):
	model = YOLO(model) # load an model
	# model.info() # display model information # only *.pt format support

	images = get_images(dir_images)
	# print("images:", images)

	os.makedirs(dir_result) #, exist_ok=True)

	for image in images:
		if torch.cuda.is_available():
			results = model.predict(dir_images+"/"+image, verbose=True, device="cuda")
		else:
			results = model.predict(dir_images+"/"+image, verbose=True)
		for result in results:
			# print(result)
			result.save(dir_result+"/"+image)
			
if __name__ == "__main__":
	# python test_yolov8_predict.py --model runs/detect/train10/weights/best_int8.engine --dir_images datasets/melon_new_detect/images/test --dir_result result_detect_engine_int8
	colorama.init()
	args = parse_args()

	if torch.cuda.is_available():
		print("Runging on GPU")
	else:
		print("Runting on CPU")

	predict(args.model, args.dir_images, args.dir_result)

	print(colorama.Fore.GREEN + "====== execution completed ======")

执行结果如下图所示：