专栏名称: AI算法与图像处理
考研逆袭985,非科班跨行AI,目前从事计算机视觉的工业和商业相关应用的工作。分享最新最前沿的科技,共同分享宝贵的资源资料,这里有机器学习,计算机视觉,Python等技术实战分享,也有考研,转行IT经验交流心得
目录
相关文章推荐
深圳大件事  ·  31名深圳业主跨市买房,都是同一“女友”找的 ... ·  15 小时前  
润农畜牧报价  ·  2025年3月13日 全国各地区鸡蛋报价! ·  昨天  
润农畜牧报价  ·  2025年3月13日 ... ·  昨天  
深圳大件事  ·  总裁辞职!深圳知名企业突发公告 ·  2 天前  
51好读  ›  专栏  ›  AI算法与图像处理

提升篇 | 优化YOLOv8以加快推理速度

AI算法与图像处理  · 公众号  ·  · 2024-11-25 22:24

正文

为了一项研究,我需要减少YOLOv8的推理时间。 在这项研究中,我使用了自己的电脑而不是Google Colab。我的电脑有一个Intel i5(第12代)处理器,我的GPU是NVIDIA GeForce RTX 3050。这些信息很重要,因为我在一些方法中使用了CPU,在其他方法中使用了GPU。

原始模型使用情况
为了测试,我们使用了Ultralytics提供的YOLOv8n.pt模型,并使用bus.jpg图像进行评估。我们将分析获得的时间值和结果。 要了解模型的性能,还要知道它运行在哪个设备上——无论是使用CUDA GPU还是CPU。
# cudaimport cv2import matplotlib.pyplot as pltfrom ultralytics import YOLOimport torch
yolov8model = YOLO("yolov8n.pt")img = cv2.imread("bus.jpg")results = yolov8model.predict(source=img, device='cuda')

img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
for result in results: boxes = result.boxes for box in boxes: x1, y1, x2, y2 = box.xyxy[0].tolist() confidence = box.conf[0].item() class_id = int(box.cls[0].item())
cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), (255, 0, 0), 2) cv2.putText(img, f'ID: {class_id} Conf: {confidence:.2f}', (int(x1), int(y1)-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
used_device = next(yolov8model.model.parameters()).deviceprint("Model is running on:", used_device)plt.figure(figsize=(10, 10))plt.imshow(img)plt.axis('off')plt.show()
# cpuimport cv2import matplotlib.pyplot as pltfrom ultralytics import YOLOimport torch
yolov8model = YOLO("yolov8n.pt")img = cv2.imread("bus.jpg")results = yolov8model.predict(source=img, device='cpu')

img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
for result in results: boxes = result.boxes for box in boxes: x1, y1, x2, y2 = box.xyxy[0].tolist() confidence = box.conf[0].item() class_id = int(box.cls[0].item())
cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), (255, 0, 0), 2) cv2.putText(img, f'ID: {class_id} Conf: {confidence:.2f}', (int(x1), int(y1)-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
plt.figure(figsize=(10, 10))plt.imshow(img)plt.axis('off')plt.show()
used_device = next(yolov8model.model.parameters()).deviceprint("Model is running on:", used_device)
现在,我们有一个起点。具体来说,对于bus.jpg图像,模型在CPU上的推理时间是199.7毫秒,在GPU上是47.2毫秒。

剪枝
我们使用的第一个方法是剪枝模型。剪枝改变了模型并创建了一个更高效的版本。有些方法修改了模型本身,而其他方法改变了输入或直接影响推理。在剪枝中,模型中较不重要或影响最小的连接被移除。这导致了一个更小、更快的模型,但它可能会对准确性产生负面影响。
import torchimport torch.nn.utils.prune as prunefrom ultralytics import YOLO
def prune_model(model,amount=0.3): for module in model.modules(): if isinstance(module,torch.nn.Conv2d): prune.l1_unstructured(module,name="weight",amount=amount) prune.remove(module,"weight") return model
model = YOLO("yolov8n.pt")#results= model.val(data="coco.yaml")
#print(f"mAP50-95: {results.box.map}")torch_model = model.modelprint(torch_model)
print("Prunning model...")pruned_torch_model = prune_model(torch_model,amount=0.1)print("Model pruned.")
model.model =pruned_torch_model
print("Saving pruned model...")model.save("yolov8n_trained_pruned.pt")
print("Pruned model saved.")
通常,一种方法被用来比较数据集;然而,在这个例子中,使用了大约18 GB的数据集的通用yolov8n.pt模型。在这个例子中,没有使用coco.yaml文件。
我将分享使用的GPU的结果,我们将更新比较图,因为应用不同的参数时时间可能会改变。通常,我无法弄清楚时间为何会改变,但这可能是由于内存或其他因素。
# cuda prunedimport cv2import matplotlib.pyplot as pltfrom ultralytics import YOLOimport torch
yolov8model = YOLO("yolov8n_trained_pruned.pt")img = cv2.imread("bus.jpg")results = yolov8model.predict(source=img, device='cuda')

img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
for result in results: boxes = result.boxes for box in boxes: x1, y1, x2, y2 = box.xyxy[0].tolist() confidence = box.conf[0].item() class_id = int(box.cls[0].item())
cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), (255, 0, 0), 2) cv2.putText(img, f'ID: {class_id} Conf: {confidence:.2f}', (int(x1), int(y1)-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2) used_device = next(yolov8model.model.parameters()).deviceprint("Model is running on:", used_device)plt.figure(figsize=(10, 10))plt.imshow(img)plt.axis('off')plt.show()
正如你看到的,结果有点令人困惑;ID和blob不准确。
然而,当我们比较推理时间时,剪枝模型在CPU和GPU上都比原始模型表现略好。剪枝模型的问题是它会影响结果,但它减少了模型的推理时间。

改变批量大小
在确定模型训练或预测的批量大小时,我们模型中同时处理的帧数至关重要。我创建了一个循环来识别最优批量大小,因为增加批量大小有时可能会产生负面影响。然而,我注意到每次尝试时最优批量大小都会改变。我尝试平均结果,但这种方法是不充分的。






请到「今天看啥」查看全文