概述

    这里主要记录使用NVIDIA GPU + pytorch + 检测系列模型的快速使用方式(包括：训练、测试、导出、量化)，可以快速解决一些工业应用的问题，比如：无网、数据大需要改路径、需要记录不同实验结果等问题。

安装

    参考官网，自己安装好Python >= 3.8和pytorch >= 1.8 后，只需要一行指令：

    pip install ultralytics

    如果要导出非pytorch版本需要安装

 pip onnx onnxslim onnxruntime

    如果需要导出TensorRT (需要注意环境变量生效的范围，如果报错“libnvinfer.so.10: cannot open shared object file: No such file or directory”)

    参考：Installation Guide :: NVIDIA Deep Learning TensorRT Documentation    

cd /home/zhangjianhu7/codes_train-test/
tar -xzvf TensorRT-10.7.0.23.Linux.x86_64-gnu.cuda-12.6.tar.gz
export LD_LIBRARY_PATH=/home/zhangjianhu7/codes_train-test/TensorRT-10.7.0.23/lib:$LD_LIBRARY_PATH  
pip install tensorrt-10.7.0-cp310-none-linux_x86_64.whl

    note

    1、对于没有外网或者想要自己掌握权重存储路径的，需要自己下载权重到你想要的位置，下载链接如下，并且选择对应的权重进行下载（目前20241202，yolov11n是必须下载的其他的都需要里面的一些权重？）：

    https://github.com/ultralytics/ultralytics?tab=readme-ov-file#python

    2、修改三个默认路径。主要是默认权重、数据、训练推理结果的读取和存储路径：Quickstart - Ultralytics YOLO Docs。修改方式可以通过命令行或者再Python代码中实时修改：Quickstart - Ultralytics YOLO Docs

训练

    我这里推荐使用data.yaml + default.yaml 修改的方式，当然也可以不使用这种方式而是通过在训练语句中传入指定参数，各有利弊。

    Data.yaml 可以参考: https://github.com/ultralytics/ultralytics/blob/1ce3c40cb3ec6ddec84d615ebe8734a33048f458/ultralytics/cfg/datasets/coco8.yaml

    default.yaml 在默认的基础上进行修改，主要修改下面的选项: https://github.com/ultralytics/ultralytics/blob/1ce3c40cb3ec6ddec84d615ebe8734a33048f458/ultralytics/cfg/default.yaml

    其他配置可以直接参考：Train - Ultralytics YOLO Docs

   

    以下是我常用的训练脚本：

from ultralytics import YOLO
from ultralytics import settings


# Update multiple settings
# settings.reset()
# value = settings["runs_dir"]
settings.update({"runs_dir": "/home/zhangjianhu/codes_train-test/yolov11",

                 "weights_dir": "/home/zhangjianhu/codes_train-test/yolov11/weights_official_8.3.32",

                "datasets_dir": "/home/zhangjianhu/datasets_project"})
print(settings)

model = YOLO(model="yolo11m.pt")
results = model.train(data="/home/zhangjianhu/codes_train-test/yolov11/jdr_xxx_241119v1/configs/datasets.yaml",
                      cfg="/home/zhangjianhu/codes_train-test/yolov11/jdr_xxx_241119v1/configs/default.yaml",

                      device=[0])

​

推理

基本等同上面的训练，大家参考官方的配置超参数就可以：Predict - Ultralytics YOLO Docs

我自己常用的代码如下：

from ultralytics import YOLO

from ultralytics import settings


# Update multiple settings
# settings.reset()
# value = settings["runs_dir"]
settings.update({"runs_dir": "/home/zhangjianhu/codes_train-test/yolov11",

                 "weights_dir": "/home/zhangjianhu/codes_train-test/yolov11/weights_official_8.3.32",

                "datasets_dir": "/home/zhangjianhu/datasets_project"})
print(settings)


model = YOLO(model="/home/zhangjianhu/codes_train-test/yolov11/jdr_xxx_241119v1/train/weights/best.pt")
results = model.predict(source="/home/zhangjianhu/datasets_project/aesthetic/jdr_xxx_241015-3w-detect/images/val",project="/home/zhangjianhu/codes_train-test/yolov11/jdr_xxx_241119v1/predict",name="jdr_xxx_241015-3w-detect-val",save=True,device=[0])

导出

        这里作者支持了直接导出TRT（并且官方提供了TRT的一系列导出教程），是直接支持TRT 的int8（PTQ）、half、dynamic：

Export - Ultralytics YOLO Docs

Export - Ultralytics YOLO Docs

TensorRT - Ultralytics YOLO Docs

导出到pytorch

from ultralytics import YOLO


model=YOLO("/home/zhangjianhu7/codes_train-test/yolov11/jdr_xxxx_241119v1/train/weights/best.pt")
model=model.cuda()
model.export(format="torchscript",device="cuda")

导出到TensorRT——FP32

from ultralytics import YOLO
model=YOLO("/home/zhangjianhu7/codes_train-test/yolov11/jdr_beauty_factor_241119v1/train/weights/best.pt")
model=model.cuda()
# model.export(format="engine",device="cuda", half=True,dynamic=True, simplify=True, batch=8)
model.export(format="engine",device="cuda")

导出到TensorRT——FP16

        这个简单但是需要先查看硬件是否支持，Tesla之后的基本都支持比如T4、A100等，老卡比如P40不支持。模型精度基本不损失，模型大小变为原来1/2，速度提升0.5-1X（需要自己测试）。

from ultralytics import YOLO

# Load the YOLOv8 model
model=YOLO("/home/zhangjianhu7/codes_train-test/yolov11/jdr_beauty_factor_241119v1/train/weights/best.pt")
model=model.cuda()

# Export the model to TensorRT format
# TensorRT FP16
model.export(format="engine", imgsz=640, dynamic=True, verbose=False, batch=8, workspace=2, half=True)

# Load the exported TensorRT model
# tensorrt_model = YOLO("/home/zhangjianhu7/codes_train-test/yolov11/jdr_beauty_factor_241119v1/train/weights/best.engine")

# Run inference
# results = tensorrt_model("https://ultralytics.com/images/bus.jpg")

测试时间

        备注：我理解这个时间不准，因为还有预处理，理论应该刨除这部分，仅仅测试模型GPU推理部分。

import cv2
from ultralytics import YOLO

model = YOLO("yolov8n.engine")
img = cv2.imread("path/to/image.jpg")

for _ in range(100):
    result = model.predict([img] * 8, verbose=False, device="cuda")

测试精度

        一般INT8有损失，fp16无损失。

from ultralytics import YOLO

model = YOLO("yolov8n.engine")
results = model.val(
    data="data.yaml",  # COCO, ImageNet, or DOTAv1 for appropriate model task
    batch=1,
    imgsz=640,
    verbose=False,
    device="cuda",
)

打赏

你的打赏是我不断分享的动力，羞羞。点这里，嘿嘿。

参考链接

Quickstart - Ultralytics YOLO Docs