【深度学习】yolov8n模型的剪枝操作记录

时间:2024-10-09 11:36:59

在这里插入图片描述
原始
在这里插入图片描述
剪枝微调后
可以看到模型大小了, 测试结果显示再cpu 上加速5%-10%

from ultralytics import YOLO

import time
# Load a pretrained YOLO11n model
count_num = 500
def test1():
    
    model = YOLO("/home/justin/Desktop/code/v8_prun/runs/detect/train3/weights/best.pt")

    # Define path to directory containing images and videos for inference
    source_img = "/home/data_for_ai_justin3/justin/JersyNumber/num_from_online_from_0822_to_0905/117.jpg"

    # # Run inference on the source
    # results = model(source, stream=True,device="cpu")  # generator of Results objects

    start_time = time.time()
    for i in range(count_num):
        model.predict(source_img, imgsz=320, conf=0.5,device="cpu")
    end_time = time.time()
    print("Time taken for inference:", (end_time - start_time)/count_num)
    return (end_time - start_time)/count_num
def test2():
    model = YOLO("/home/justin/Desktop/code/v8_prun/runs/detect/train4/weights/best.pt")

    # Define path to directory containing images and videos for inference
    source_img = "/home/data_for_ai_justin3/justin/JersyNumber/num_from_online_from_0822_to_0905/117.jpg"

    # # Run inference on the source
    # results = model(source, stream=True,device="cpu")  # generator of Results objects
    
    start_time = time.time()
    for i in range(count_num):
        model.predict(source_img, imgsz=320, conf=0.5,device="cpu")
    end_time = time.time()
    print("Time taken for inference:", (end_time - start_time)/count_num)
    return (end_time - start_time)/count_num

if __name__ == '__main__':
    a1 = test1()
    a2 = test2()
    print("Speedup:", a1/a2,a1,a2)

测试代码
剪枝的ref:https://blog.****.net/Dneccc/article/details/134440603?utm_medium=distribute.pc_relevant.none-task-blog-2defaultbaidujs_baidulandingword~default-5-134440603-blog-130155924.235v43pc_blog_bottom_relevance_base2&spm=1001.2101.3001.4242.4&utm_relevant_index=8