![cover](https://i-blog.csdnimg.cn/direct/3a8324a80ec0488bbc9143e025734260.png)
将YOLOv8模型从PyTorch的.pt格式转换为TensorRT的.engine格式
将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 ======")
执行结果如下图所示:
预测结果图像如下所示:
更多推荐
所有评论(0)