本文將通過一個完整的實戰(zhàn)案例來展示使用Python、命令行或Google Colab等方式在自定義數(shù)據(jù)集上訓(xùn)練自己的計算機(jī)視覺模型。

簡介

當(dāng)前，Ultralytics公司研制的高級YOLOv8模型成為解決計算機(jī)視覺問題的最佳方法之一，同時該模型也最大限度地減少了有關(guān)開發(fā)過程遭遇的麻煩。YOLOv8是Ultralytics YOLO（You Only Look Once）系列模型的第8次也是最新一次迭代，與其他迭代一樣，它使用卷積神經(jīng)網(wǎng)絡(luò)（CNN）來預(yù)測對象類別及其邊界框。YOLO系列對象探測器以準(zhǔn)確和快速而聞名，并提供了一個基于PyTorch構(gòu)建的平臺，從而簡化了人們從頭開始創(chuàng)建模型的大部分過程。

尤其重要的是，YOLOv8也是一個非常靈活的模型。它可以在各種平臺上使用你選擇的任何數(shù)據(jù)集進(jìn)行訓(xùn)練，預(yù)測模型可以基于許多數(shù)據(jù)源運行。本文將作為一個全面的教程，涵蓋訓(xùn)練和運行YOLOv8模型的許多不同方法，并展示每種方法的優(yōu)缺點，這些方法將根據(jù)你的硬件和數(shù)據(jù)集幫你選擇最合適的開發(fā)方案。

【注意】本文中創(chuàng)建上述示例數(shù)據(jù)集時使用的所有圖像均由作者本人拍攝。

開發(fā)環(huán)境搭建

要開始訓(xùn)練我們的YOLOv8模型，第一步是決定我們想在什么樣的環(huán)境中訓(xùn)練我們的模型（請記住，訓(xùn)練和運行模型是單獨的任務(wù)）。

歸納來看，可供我們選擇的環(huán)境在很大程度上可以分為兩類：基于本地運行和基于云端方案運行。

通過基于本地的訓(xùn)練，我們基本上是使用設(shè)備的物理硬件直接在我們的系統(tǒng)上運行訓(xùn)練過程。在基于本地的訓(xùn)練中，YOLOv8為我們提供了兩個選項：Python API和CLI。這兩種選擇的結(jié)果或速度沒有真正的區(qū)別，因為同樣的過程是在幕后進(jìn)行的；唯一的區(qū)別在于訓(xùn)練的設(shè)置和運行方式。

另一方面，基于云端的訓(xùn)練允許你利用云服務(wù)器的硬件。通過使用互聯(lián)網(wǎng)，你可以連接到云運行時并執(zhí)行代碼，就像在本地計算機(jī)上一樣，只是現(xiàn)在它在云硬件上運行。

到目前為止，最受歡迎的機(jī)器學(xué)習(xí)云平臺是谷歌Colab。它使用Jupyter筆記本格式，允許用戶創(chuàng)建可以編寫和運行代碼片段的“單元格”，并提供與Google Drive和Github的強(qiáng)大集成。

你決定使用哪種環(huán)境在很大程度上取決于你可用的硬件。如果你有一個配備高端NVIDIA GPU的強(qiáng)大系統(tǒng)，基于本地的訓(xùn)練可能會對你很有效。如果你的本地機(jī)器的硬件不符合機(jī)器學(xué)習(xí)的配置建議，或者你只是想要比本地更多的計算能力，谷歌Colab可能是你更恰當(dāng)?shù)倪x擇。

谷歌Colab最大的好處之一是它免費提供了一些計算資源，但也有一個簡單的升級路徑，可以讓你利用更快的計算硬件。即使你已經(jīng)有了一個強(qiáng)大的系統(tǒng)，如果谷歌Colab的高層計劃中提供的更快的GPU比你現(xiàn)有的硬件有顯著的性能提升，你也可以考慮使用谷歌Colab。使用免費計劃，你只能使用NVIDIA T4，其性能大致相當(dāng)于RTX 2070。對于更高級別配置的計劃，L4（性能約為4090）和A100（性能約2 4090）是可用的。在比較GPU時，請記住VRAM的數(shù)量是機(jī)器學(xué)習(xí)性能的主要決定因素。

數(shù)據(jù)集

為了開始訓(xùn)練模型，你需要大量數(shù)據(jù)來訓(xùn)練它。對象檢測數(shù)據(jù)集通常由各種對象的圖像集合組成，此外還有一個圍繞對象的“邊界框”，指示其在圖像中的位置。

檢測到的對象周圍的邊界框示例

YOLOv8兼容的數(shù)據(jù)集具有特定的結(jié)構(gòu)。它們主要分為有效文件夾、訓(xùn)練文件夾和測試文件夾，分別用于模型的驗證、訓(xùn)練和測試（驗證和測試之間的區(qū)別在于，在驗證過程中，結(jié)果用于調(diào)整模型以提高其準(zhǔn)確性；而在測試過程中，結(jié)果僅用于提供模型真實世界準(zhǔn)確性的衡量標(biāo)準(zhǔn)）。

在每個文件夾中，數(shù)據(jù)集進(jìn)一步分為兩個文件夾：圖像（images）和標(biāo)簽（labels）文件夾。這兩個文件夾的內(nèi)容緊密相連。

顧名思義，images文件夾包含數(shù)據(jù)集的所有對象圖像。這些圖像通常具有方形縱橫比、低分辨率和小的文件尺寸。

labels文件夾包含邊界框在每個圖像中的位置和大小的數(shù)據(jù)，以及每個圖像表示的對象的類型（或類別）。例如：

這里的第一行數(shù)據(jù)代表圖像中存在的單個對象。在每一行中，第一個數(shù)字表示對象的類別，第二個和第三個數(shù)字表示邊界框中心的x和y坐標(biāo)，第四個和第五個數(shù)字表示邊緣框的寬度和高度。

images和labels文件夾中的數(shù)據(jù)通過文件名鏈接在一起。images文件夾中的每個圖像在labels文件夾中都有一個具有相同文件名的相應(yīng)文件；反之亦然。在數(shù)據(jù)集中，images和labels文件夾中總是有匹配的文件對，它們具有相同的文件名，但文件擴(kuò)展名不同；.jpg用于圖像，.txt用于標(biāo)簽。.jpg圖片中每個對象的邊界框數(shù)據(jù)包含在相應(yīng)的.txt文件中。

YOLOv8兼容數(shù)據(jù)集的典型文件結(jié)構(gòu)。來源：Ultralytics公司??YOLO文檔??

有幾種方法可以獲得與YOLOv8兼容的數(shù)據(jù)集來開始訓(xùn)練模型。你可以創(chuàng)建自己的數(shù)據(jù)集，也可以使用互聯(lián)網(wǎng)上預(yù)先配置的數(shù)據(jù)集。在本教程中，我們將使用??CVAT??創(chuàng)建自己的數(shù)據(jù)集，并使用??Kaggle??查找預(yù)配置的數(shù)據(jù)集。

CVAT

CVAT（CVAT.ai）是一個注釋工具，它允許你通過手動方式為圖像和視頻添加標(biāo)簽來創(chuàng)建自己的數(shù)據(jù)集。

創(chuàng)建帳戶并登錄后，開始注釋的過程很簡單。只需創(chuàng)建一個項目，給它一個合適的名稱，并根據(jù)需要為盡可能多的對象類型/類別添加標(biāo)簽。

作者在cvat.ai.Video上創(chuàng)建新項目和標(biāo)簽

創(chuàng)建一個新任務(wù)，并上傳你希望成為數(shù)據(jù)集一部分的所有圖像。點擊“Submit & Open”按鈕，項目下應(yīng)創(chuàng)建一個新任務(wù)，其中包含一個作業(yè)。

在cvat.ai上創(chuàng)建新任務(wù)和作業(yè)

打開此作業(yè)將允許你啟動注釋過程。你可以使用矩形工具為數(shù)據(jù)集中的每個圖像創(chuàng)建邊界框和標(biāo)簽。

使用cvat.ai上的矩形工具創(chuàng)建邊界框

在為所有圖像添加注釋后，返回任務(wù)并選擇“Actions（動作）”→“Export task dataset（導(dǎo)出任務(wù)數(shù)據(jù)集）”，然后選擇YOLOv8 Detection 1.0作為導(dǎo)出格式。下載任務(wù)數(shù)據(jù)集后，你會發(fā)現(xiàn)它只包含labels文件夾，而不包含images文件夾（除非你在導(dǎo)出時選擇了“Save images（保存圖像）”選項）。你必須手動創(chuàng)建images文件夾并將圖像移動到那里（你可能希望首先將圖像壓縮到較低的分辨率，例如640x640）。記住不要更改文件名，因為它們必須與labels文件夾中.txt文件的文件名匹配。你還需要決定如何在驗證、訓(xùn)練和測試之間分配圖像（訓(xùn)練是其中最重要的）。

作者從cvat.ai.Image導(dǎo)出的示例數(shù)據(jù)集

至此，你的數(shù)據(jù)集已完成并準(zhǔn)備好使用了！

Kaggle

Kaggle（Kaggle.com）是最大的在線數(shù)據(jù)科學(xué)社區(qū)之一，也是探索數(shù)據(jù)集的最佳網(wǎng)站之一。你可以通過簡單地搜索他們的網(wǎng)站來嘗試找到你需要的數(shù)據(jù)集，除非你正在尋找非常具體的東西；否則，你很可能會找到意向的數(shù)據(jù)集。然而，Kaggle上的許多數(shù)據(jù)集不是YOLOv8兼容的格式和/或與計算機(jī)視覺無關(guān)；因此，你可能想在查詢中包含“YOLOv9”來優(yōu)化你的搜索。

你可以通過數(shù)據(jù)集的Data Explorer（頁面右側(cè)）中的文件結(jié)構(gòu)來判斷數(shù)據(jù)集是否與YOLOv8兼容。

Kaggle上兼容YOLOv8的數(shù)據(jù)集示例

如果數(shù)據(jù)集相對較小（幾MB）和/或你想在本地訓(xùn)練，那么可以直接從Kaggle下載數(shù)據(jù)集。但是，如果你計劃在Google Colab上使用大型數(shù)據(jù)集進(jìn)行訓(xùn)練，最好從筆記本文件本身進(jìn)行數(shù)據(jù)集檢索（更多信息見下文）。

訓(xùn)練模型

訓(xùn)練過程將根據(jù)你是在本地還是在云端進(jìn)行訓(xùn)練而有所不同。

本地

為所有訓(xùn)練文件創(chuàng)建一個項目文件夾。在本教程中，我們將稱之為yolov8-project。然后，將數(shù)據(jù)集移動/復(fù)制到此文件夾下。

接下來，使用所需的YOLOv8依賴項設(shè)置Python虛擬環(huán)境：

python3 -m venv venv
source venv/bin/activate
pip3 install ultralytics

然后，創(chuàng)建一個名為config.yaml的配置文件。這個文件將會指定用于訓(xùn)練的重要數(shù)據(jù)集信息：

path: /Users/oliverma/yolov8-project/dataset/ # absolute path to dataset
test: test/images # relative path to test images
train: train/images # relative path to training images
val: val/images # relative path to validation images

# classes
names:
0: bottle

在上面配置信息的path部分，提供的是指向數(shù)據(jù)集根目錄的絕對文件路徑。你也可以使用相對文件路徑，但這取決于config.yaml的相對位置。

然后，在test、train和val部分，提供用于測試、訓(xùn)練和驗證的圖像的位置（如果你只有訓(xùn)練圖像，則只需對所有這三種操作均使用train/images）。

在names部分，指定每個類別的名稱。這些信息通常可以在任何YOLOv8數(shù)據(jù)集的data.yaml文件中找到。

如前所述，Python API或CLI（命令行方式）都可以用來進(jìn)行本地訓(xùn)練。

Python API方式

創(chuàng)建另一個名為main.py的文件。這是實際訓(xùn)練開始的地方：

from ultralytics import YOLO
model = YOLO("yolov8n.yaml")
model.train(data="config.yaml", epochs=100)

通過將我們的模型初始化為YOLO("yolov8n.yaml")，我們基本上是從頭開始創(chuàng)建一個新模型。我們使用yolov8n是因為它是最快的模型，但根據(jù)你自己的使用情況，你也可以選擇使用其他模型。

性能指標(biāo)

YOLOv8變體的性能指標(biāo)。來源：??Ultralytics YOLO文檔??

最后，我們開始訓(xùn)練模型，并傳遞配置文件和迭代次數(shù)，或訓(xùn)練輪數(shù)。一個比較好的訓(xùn)練指標(biāo)是使用300個訓(xùn)練輪數(shù)，但你可能想根據(jù)數(shù)據(jù)集的大小和硬件的速度來調(diào)整這個數(shù)字。

你可能還希望包括一些更有用的設(shè)置：

• imgsz：將所有圖像調(diào)整到指定的大小。例如，imgsz=640會將所有圖像的大小調(diào)整為640x640。如果你創(chuàng)建了自己的數(shù)據(jù)集并且沒有調(diào)整圖像大小，這將非常有用。
• device：指定要在哪個設(shè)備上訓(xùn)練。默認(rèn)情況下，YOLOv8會嘗試在GPU上訓(xùn)練，并使用CPU訓(xùn)練作為后備，但如果你在M系列Mac上訓(xùn)練，你必須使用device="mps"以便使用蘋果電腦上的Metal Performance Shaders（mps）后端進(jìn)行GPU加速訓(xùn)練。
• 有關(guān)所有訓(xùn)練參數(shù)的更多信息，請訪問https://docs.ultralytics.com/modes/train/#train-settings。

你的項目目錄現(xiàn)在應(yīng)該看起來類似于下面的樣子：

項目目錄的示例文件結(jié)構(gòu)

現(xiàn)在，我們終于準(zhǔn)備好開始訓(xùn)練我們的模型了。為此，只需要在項目目錄中打開一個終端并運行：

python3 main.py

隨著訓(xùn)練的進(jìn)行，終端將顯示每個訓(xùn)練世代的訓(xùn)練進(jìn)度信息。

終端中顯示的每個訓(xùn)練世代的訓(xùn)練進(jìn)度

訓(xùn)練結(jié)果將保存在路徑runs/detect/train（或者train2，train3，等）。注意，這里包括了權(quán)重數(shù)據(jù)（文件擴(kuò)展名為.pt），這對以后運行模型很重要；還有文件results.png，它顯示了許多包含相關(guān)訓(xùn)練統(tǒng)計數(shù)據(jù)的圖表。

results.png文件中的示例圖

CLI方式

在項目目錄中打開一個新終端并運行以下命令：

yolo detect train data=config.yaml model=yolov8n.yaml epochs=100

此命令可以使用上面為Python API列出的相同參數(shù)進(jìn)行修改。例如：

yolo detect train data=config.yaml model=yolov8n.yaml epochs=300 imgsz=640 device=mps

訓(xùn)練將開始，進(jìn)度將顯示在終端上。其余的訓(xùn)練過程與Python CLI相同。

谷歌Colab方式

導(dǎo)航到https://colab.research.google.com/，并為訓(xùn)練創(chuàng)建一個新的筆記本文件。

在訓(xùn)練之前，請確保通過選擇右上角的“Change runtime type（更改運行時類型）”連接到GPU運行時。CPU運行時的訓(xùn)練將非常緩慢。

將筆記本電腦運行時從CPU更改為T4 GPU

在開始使用Google Colab進(jìn)行任何訓(xùn)練之前，我們首先需要將數(shù)據(jù)集導(dǎo)入筆記本文件。直觀地說，最簡單的方法是將數(shù)據(jù)集上傳到谷歌云端硬盤，并從那里導(dǎo)入到我們的筆記本文件中。然而，上傳任何大于幾MB的數(shù)據(jù)集都需要非常長的時間。解決方法是將數(shù)據(jù)集上傳到遠(yuǎn)程文件托管服務(wù)（如Amazon S3甚至是Kaggle），并將數(shù)據(jù)集直接從那里拉入我們的Colab筆記本文件。

從Kaggle導(dǎo)入

以下是如何將Kaggle數(shù)據(jù)集直接導(dǎo)入Colab筆記本文件的說明：

在Kaggle賬戶設(shè)置中，向下滾動到API并選擇“Create New Token”命令。這將下載一個名為kaggle.json的文件。

在筆記本單元格中運行以下命令：

!pip install kaggle
from google.colab import files
files.upload()

上傳剛剛下載的kaggle.json文件，然后運行以下命令：

!mkdir ~/.kaggle
!cp kaggle.json ~/.kaggle/
!chmod 600 ~/.kaggle/kaggle.json
!kaggle datasets download -d [DATASET] # replace [DATASET] with the desired dataset ref

數(shù)據(jù)集將作為zip存檔下載。只需要使用unzip命令提取有關(guān)內(nèi)容：

!unzip dataset.zip -d dataset

開始訓(xùn)練

在筆記本文件的文件資源管理器中創(chuàng)建一個新的config.yaml文件，并按照前面的描述進(jìn)行配置。Colab筆記本文件中的默認(rèn)工作目錄是/content/，因此數(shù)據(jù)集的絕對路徑將是/content/[dataset folder]。例如：

path: /content/dataset/ # absolute path to dataset
test: test/images # relative path to test images
train: train/images # relative path to training images
val: val/images # relative path to validation images

# classes
names:
0: bottle

確保檢查一下數(shù)據(jù)集的文件結(jié)構(gòu)，以確保config.yaml中指定的路徑準(zhǔn)確。有時數(shù)據(jù)集會被放置在多個級別的文件夾中。

然后，將以下內(nèi)容作為單元格運行：

!pip install ultralytics

import os
from ultralytics import YOLOmodel = YOLO("yolov8n.yaml")
results = model.train(data="config.yaml", epochs=100)

前面提到的用于修改本地訓(xùn)練設(shè)置的參數(shù)也適用于這里。

與本地訓(xùn)練類似，結(jié)果、權(quán)重和圖表將保存在runs/detect/train中。

在本地進(jìn)行預(yù)測

無論你是在本地還是在云端進(jìn)行訓(xùn)練，預(yù)測都必須在本地運行。

在模型完成訓(xùn)練后，runs/detect/train/weights中會有兩個權(quán)重，分別命名為best.pt和last.pt，分別是最佳訓(xùn)練輪次和最新訓(xùn)練輪次的權(quán)重。在本教程中，我們將使用best.pt運行模型。

如果你在本地進(jìn)行訓(xùn)練，請將best.pt移動到方便的位置（例如我們的項目文件夾yolov8-project）以運行預(yù)測。如果你在云端訓(xùn)練，請將best.pt下載到你的本地設(shè)備上。在Google Colab上，右鍵單擊筆記本資源管理器中的文件，然后選擇“Download”（下載）命令。

在谷歌Colab上下載權(quán)重

與本地訓(xùn)練類似，預(yù)測可以通過Python API或CLI運行。

使用Python API預(yù)測

在與best.pt相同的位置，創(chuàng)建一個名為predict.py的新文件：

from ultralytics import YOLO
model = YOLO("best.pt")
results = model(source=0, show=True, conf=0.25, save=True)

與訓(xùn)練類似，存在許多有用的參數(shù)可以修改預(yù)測設(shè)置：

• source：控制預(yù)測的輸入源。source=0將網(wǎng)絡(luò)攝像頭設(shè)置為輸入源。更多信息請見下文。
• show：如果為True，則在屏幕上顯示預(yù)測、邊界框和置信度。
• conf：要考慮預(yù)測的最小置信度閾值。
• save：如果為True，則將預(yù)測結(jié)果保存到run/retect/repredict（或predict2、predict3等）。
• device：如前所述，在M系列Mac上使用device=“mps”。

有關(guān)預(yù)測參數(shù)的完整列表，請訪問：

??https://docs.ultralytics.com/modes/predict/#inference-arguments??。

使用CLI預(yù)測

運行以下CLI命令即可啟動模型：

python3 predict.py

通過實時網(wǎng)絡(luò)攝像頭饋送運行YOLOv8模型預(yù)測

CLI命令

使用CLI方式進(jìn)行預(yù)測應(yīng)用的命令是：

yolo detect predict model=best.pt source=0 show=True conf=0.25 save=True

注意到，這里所使用的參數(shù)與Python API中的參數(shù)相同。

將YOLOv8模型集成到實際場景

現(xiàn)在，我們已經(jīng)能夠在實時網(wǎng)絡(luò)攝像頭上成功地運行起我們的模型，但這又會怎樣呢？我們?nèi)绾螌嶋H使用這個模型并將其集成到實戰(zhàn)項目中？

讓我們從輸入和輸出的角度來考慮這個問題。為了使這個模型成功地應(yīng)用于我們構(gòu)建的外部應(yīng)用程序中，此模型必須能夠接受有用的輸入并產(chǎn)生有用的輸出。值得慶幸的是，YOLOv8模型的靈活性使得將模型集成到各種使用場景中成為可能。

例如，我們可以使用source=0將網(wǎng)絡(luò)攝像頭設(shè)置為預(yù)測的輸入源。然而，YOLOv8模型可以利用比這更多的輸入源。以下是幾個例子：

results = model(source="path/to/image.jpg", show=True, conf=0.25, save=True) # static image
results = model(source="screen", show=True, conf=0.25, save=True) # screenshot of current screen
results = model(source="https://ultralytics.com/images/bus.jpg", show=True, conf=0.25, save=True) # image or video URL
results = model(source="path/to/file.csv", show=True, conf=0.25, save=True) # CSV file
results = model(source="path/to/video.mp4", show=True, conf=0.25, save=True) # video file
results = model(source="path/to/dir", show=True, conf=0.25, save=True) # all images and videos within directory
results = model(source="path/to/dir/**/*.jpg", show=True, conf=0.25, save=True) # glob expression
results = model(source="https://www.youtube.com/watch?v=dQw4w9WgXcQ", show=True, conf=0.25, save=True) # YouTube video URL

有關(guān)預(yù)測源和輸入選項的完整列表，請訪問：

??https://docs.ultralytics.com/modes/predict/#inference-sources。?

每當(dāng)我們運行預(yù)測時，YOLOv8都會以Results對象列表的形式返回大量有價值的數(shù)據(jù)，其中包括有關(guān)預(yù)測的邊界框、分割掩碼、關(guān)鍵點、類別概率和定向包圍盒（OBB）的信息。

由于我們在代碼中將預(yù)測結(jié)果分配給results變量，因此我們可以使用它來檢索有關(guān)預(yù)測的信息：

from ultralytics import YOLO

model = YOLO("best.pt")

results = model(source="bottles.jpg", show=True, conf=0.25, save=True)

print("Bounding boxes of all detected objects in xyxy format:")
for r in results:
print(r.boxes.xyxy)

print("Confidence values of all detected objects:")
for r in results:
print(r.boxes.conf)

print("Class values of all detected objects:")
for r in results:
print(r.boxes.cls)

本教程中包含的輸出結(jié)果類型太多，但你可以通過訪問鏈接以了解更多信息：

??https://docs.ultralytics.com/modes/predict/#working-results。?

這里提供的只是一個非常基本的例子，說明你可以用YOLOv8模型的輸出做點什么了。其實，你可以通過上述列舉的很多種方式將模型應(yīng)用于自己的項目中。

結(jié)論

在本文中，我們終于實現(xiàn)了從頭開始制作我們自己的YOLOv8兼容數(shù)據(jù)集，從Kaggle導(dǎo)入數(shù)據(jù)集，使用包括Python API、CLI和Google Colab在內(nèi)的多種環(huán)境來訓(xùn)練模型，然后在本地運行我們的模型，并找到許多輸入/輸出方法，使我們能夠在自己的項目中利用訓(xùn)練出的YOLOv8模型。

請記住，本教程的目的是作為YOLOv8或計算機(jī)視覺的學(xué)習(xí)起點。我們當(dāng)前幾乎沒有觸及YOLOv8模型錯綜復(fù)雜的表面，隨著你對YOLOv9和計算機(jī)視覺越來越有經(jīng)驗，深入了解這個模型絕對是明智的。

話雖如此，如果你遵循了本教程并堅持到最后，那仍然是一個很大的成功。我希望這篇文章能幫助你對機(jī)器學(xué)習(xí)、計算機(jī)視覺和YOLOv8模型有一個基本的了解。也許你甚至對這個主題產(chǎn)生了熱情，并將在未來繼續(xù)學(xué)習(xí)更高級的內(nèi)容。

譯者介紹

朱先忠，51CTO社區(qū)編輯，51CTO專家博客、講師，濰坊一所高校計算機(jī)教師，自由編程界老兵一枚。

原文標(biāo)題：??The Comprehensive Guide to Training and Running YOLOv8 Models on Custom Datasets??，作者：Oliver Ma