YOLOv8的FPS计算代码

最新推荐文章于 2024-05-07 15:21:57 发布

王清江唷_全网同名

最新推荐文章于 2024-05-07 15:21:57 发布

阅读量1.2k

点赞数 4

文章标签： YOLO 人工智能机器学习

本文链接：https://blog.csdn.net/qq_29762001/article/details/136997098

版权

本文提供了一个Python代码片段，展示了如何在YOLOv8模型上进行FPS计算，包括如何将模型加载到指定GPU，以及如何通过预热和检测一定数量的图像来获取更稳定的推理速度。

摘要由CSDN通过智能技术生成

YOLOv8的FPS计算代码

目前是默认加载到0号GPU中，如果你想加载到指定GPU中，请手动在加载模型的时候设置
device=编号

代码

import os

from ultralytics import YOLO


def load_model(model_path):
    model = YOLO(model_path)
    print('查看当前模型：', model)
    return model


if __name__ == '__main__':
    imgs_path = 'E:\\data\\dataset\\images\\val'  # TODO 验证集目录   要求图像数量 >= 200
    model = load_model('weights/shwd/yuan.pt' )  # TODO 模型路径
    re_num = 10  # TODO 预热图像张数
    detect_count = 100  # TODO 推理图像张数
    images = os.listdir(imgs_path)
    count = 0
    times = []
    for item in images:
        if count < re_num:
            model(imgs_path + os.sep + item)
        elif count < re_num + detect_count:
            results = model(imgs_path + os.sep + item)
            times.append(results[0].speed)
        else:
            break
        count += 1
    # main(opt)
    time_sum = 0.
    for i in times:
        time_sum += sum(i.values())
    one_img_time = time_sum / detect_count
    FPS = 1000 / one_img_time
    print(f'FPS: {FPS}')

代码说明

以上的代码中，需要自己改的参数有四个

    imgs_path = 'E:\\data\\dataset\\images\\val'  # TODO 验证集目录   要求图像数量 >= 200
    model = load_model('weights/shwd/yuan.pt' )  # TODO 模型路径
    re_num = 10  # TODO 预热图像张数
    detect_count = 100  # TODO 推理图像张数