使用回調(diào)函數(shù)訓(xùn)練YOLO模型
大多數(shù)人可能熟悉如何訓(xùn)練計(jì)算機(jī)視覺(jué)模型，比如流行的YOLO模型，甚至知道如何使用這些模型進(jìn)行預(yù)測(cè)。但你知道我們可以通過(guò)回調(diào)函數(shù)為這些模型增加一些靈活性，以便在模型訓(xùn)練和模型推斷中使用嗎？大多數(shù)最先進(jìn)的（SOTA）YOLO模型，如YOLOv8和YOLO-NAS，都實(shí)現(xiàn)了回調(diào)函數(shù)，我們可以調(diào)整這些函數(shù)以有效地利用我們的計(jì)算機(jī)視覺(jué)模型的訓(xùn)練和推斷。
考慮以下情景。假設(shè)你是一名計(jì)算機(jī)視覺(jué)工程師，與團(tuán)隊(duì)中的許多工程師一起工作。你正在使用自定義數(shù)據(jù)集訓(xùn)練自定義的計(jì)算機(jī)視覺(jué)模型（也許是YOLO），以實(shí)現(xiàn)一些業(yè)務(wù)邏輯。你負(fù)責(zé)實(shí)現(xiàn)訓(xùn)練和推斷邏輯。除此之外，你還需要報(bào)告模型的訓(xùn)練進(jìn)度、訓(xùn)練模型的準(zhǔn)確性等。作為一名工程師，你決定在很多個(gè)epoch上訓(xùn)練你的模型，這可能需要幾天的時(shí)間，具體取決于一些因素，比如數(shù)據(jù)集的數(shù)量、服務(wù)器資源等。你需要密切關(guān)注模型的訓(xùn)練進(jìn)度，因?yàn)橛捎谥T如服務(wù)器資源問(wèn)題等原因，模型可能在一段時(shí)間后停止訓(xùn)練，導(dǎo)致訓(xùn)練崩潰。你可能也希望在模型訓(xùn)練完成后收到自動(dòng)警報(bào)，比如在訓(xùn)練結(jié)束后收到帶有驗(yàn)證指標(biāo)的電子郵件，或者在模型訓(xùn)練完成后自動(dòng)向團(tuán)隊(duì)負(fù)責(zé)人發(fā)送報(bào)告。這些以及許多其他事情都是你作為計(jì)算機(jī)視覺(jué)工程師可能想要做的事情。
要實(shí)現(xiàn)以上任何一種情況，我們需要一種回調(diào)函數(shù)。這就是在訓(xùn)練計(jì)算機(jī)視覺(jué)模型時(shí)回調(diào)函數(shù)的作用。好消息是，大多數(shù)SOTA YOLO模型默認(rèn)實(shí)現(xiàn)了這些回調(diào)函數(shù)。例如，默認(rèn)情況下，YOLOv8和YOLO-NAS實(shí)現(xiàn)了這些回調(diào)函數(shù)，你可以在訓(xùn)練或進(jìn)行模型預(yù)測(cè)時(shí)有效地利用它們。在本文章中，我將向你展示一些示例，演示在訓(xùn)練YOLO模型時(shí)如何使用回調(diào)函數(shù)。在本例中，我將使用YOLOv8，但請(qǐng)注意，這可以擴(kuò)展到其他一些YOLO模型，比如YOLO-NAS。
讓我們繼續(xù)演示如何在YOLOv8上實(shí)現(xiàn)回調(diào)函數(shù)。我們將編寫(xiě)代碼并在自定義數(shù)據(jù)集上訓(xùn)練我們的模型。我們將實(shí)現(xiàn)回調(diào)函數(shù)。其中一個(gè)功能是在模型訓(xùn)練結(jié)束后向我們的團(tuán)隊(duì)工程師發(fā)送電子郵件。我們發(fā)送的電子郵件將包含受過(guò)訓(xùn)練模型的報(bào)告，如指標(biāo)、訓(xùn)練模型所花費(fèi)的時(shí)間等。
項(xiàng)目實(shí)施步驟
第1步：創(chuàng)建一個(gè)文件夾并給它命名（在我的案例中，我將我的文件夾命名為“yolo_with_callbacks”）。
在你創(chuàng)建的文件夾中，創(chuàng)建一個(gè)新的文本文件（requirements.txt）并添加以下內(nèi)容：
opencv-python==4.8.1.78
Pillow==10.0.1
tqdm==4.66.1
ultralytics==8.1.2
python-dotenv==1.0.1然后，在你的項(xiàng)目文件夾中創(chuàng)建一個(gè)Python虛擬環(huán)境，并安裝requirements.txt文件中列出的依賴項(xiàng)。
python3 -m venv env接下來(lái)，通過(guò)運(yùn)行以下命令激活新創(chuàng)建的虛擬環(huán)境：
source env/bin/activate # if you are using Ubuntu
source env/Scripts/activate # if you are using Windows然后，通過(guò)運(yùn)行以下命令安裝依賴項(xiàng)：
pip install -r requirements.txt第2步：下載一個(gè)用于自定義模型訓(xùn)練的示例數(shù)據(jù)集。
你可以使用任何你選擇的數(shù)據(jù)集，只要注釋是以YOLO格式提供的即可。在我的案例中，為了本教程的目的，我將使用來(lái)自Roboflow的POTHOLE數(shù)據(jù)集，你可以從這個(gè)鏈接下載：POTHOLE數(shù)據(jù)集。下載數(shù)據(jù)集后，你將得到三個(gè)文件夾（train、val和test）?，F(xiàn)在，在你的項(xiàng)目目錄中創(chuàng)建一個(gè)數(shù)據(jù)集文件夾，并將你下載的數(shù)據(jù)集（train、val和test）復(fù)制到這個(gè)文件夾中。你的數(shù)據(jù)集文件夾應(yīng)該如下所示：
Datasets
└── train
├── images
└── labels
└── val
├── images
└── labels
└── test
├── images
└── labels接下來(lái)，在項(xiàng)目根目錄中創(chuàng)建一個(gè)數(shù)據(jù)集配置文件（我們稱之為data.yaml）并在YAML文件中添加以下內(nèi)容：
train: ./dataset/train/images
val: ./dataset/val/images
test: ./dataset/test/images
nc: 1
names: ['pothole']第3步：創(chuàng)建模型訓(xùn)練腳本。
接下來(lái)，我們需要編寫(xiě)代碼來(lái)使用我們的自定義數(shù)據(jù)集訓(xùn)練模型。之后，我們將繼續(xù)實(shí)現(xiàn)模型的回調(diào)函數(shù)，這是本教程的唯一目的?，F(xiàn)在，在你的項(xiàng)目根目錄中創(chuàng)建一個(gè)新文件（命名為training.py）。在這個(gè)training.py文件中，我們將實(shí)現(xiàn)模型訓(xùn)練和回調(diào)函數(shù)。首先，讓我們編寫(xiě)一個(gè)用于訓(xùn)練YOLOV8模型的函數(shù)：
def train_yolov8_model(config_path, num_epochs, training_result_dir):
model = YOLO("yolov8x.pt")
model.add_callback("on_train_start", on_train_start)
model.add_callback("on_train_epoch_end", on_train_epoch_end)
model.add_callback("on_train_end", on_train_end)
model.start_time = datetime.now()
start_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
# Train the model
model.train(
data=config_path,
name="Yolo_Model_Training",
project=training_result_dir,
task="detect",
epochs=num_epochs,
patience=20,
batch=16,
cache=True,
imgsz=640,
iou=0.5,
augment=True,
degrees=25.0,
fliplr=0.0,
lr0=0.0001,
optimizer="Adam",
device=device,
)注意：函數(shù)參數(shù)中的config_path是我們之前創(chuàng)建的數(shù)據(jù)集yaml配置文件。我們稍后將定義的回調(diào)函數(shù)，就像model.add_callback這樣的調(diào)用，稍等一下。
接下來(lái)，讓我們實(shí)現(xiàn)回調(diào)函數(shù)。在這種情況下，我們將要實(shí)現(xiàn)的回調(diào)函數(shù)包括：on_train_start、on_train_epoch_end和on_train_end。on_train_start回調(diào)是在模型開(kāi)始訓(xùn)練時(shí)立即觸發(fā)的回調(diào)函數(shù)。on_train_epoch_end是在每個(gè)epoch結(jié)束后立即觸發(fā)的回調(diào)函數(shù)。on_train_end是在模型完成訓(xùn)練后觸發(fā)的回調(diào)函數(shù)。
實(shí)現(xiàn)回調(diào)函數(shù)
def on_train_start(trainer):
start_time = datetime.now()
def on_train_epoch_end(trainer):
curr_epoch = trainer.epoch + 1
text = f"Epoch Number: {curr_epoch}/{trainer.epochs} finished"
print(text)
print("-" * 50)對(duì)于on_train_start回調(diào)，我們需要追蹤模型開(kāi)始訓(xùn)練的確切時(shí)間。你實(shí)際上可以在這里實(shí)現(xiàn)更復(fù)雜的邏輯。對(duì)于on_train_epoch_end，我們只是獲取了當(dāng)前epoch并打印出來(lái)。這只是一個(gè)簡(jiǎn)單的演示。我們可以在這里實(shí)現(xiàn)更復(fù)雜的邏輯。例如，如果我們有一個(gè)用戶正在從中訓(xùn)練模型的前端應(yīng)用程序，我們可以在每個(gè)epoch結(jié)束后更新GUI的訓(xùn)練進(jìn)度條。我們可以在這個(gè)函數(shù)中實(shí)現(xiàn)這個(gè)功能。
現(xiàn)在，讓我們繼續(xù)實(shí)現(xiàn)本教程的主要邏輯。我們將繼續(xù)實(shí)現(xiàn)on_train_end回調(diào)函數(shù)。如前所述，此函數(shù)僅在模型訓(xùn)練成功完成后觸發(fā)。在我們的情況下，我們想要發(fā)送一個(gè)包含模型訓(xùn)練報(bào)告的電子郵件給我們的團(tuán)隊(duì)工程師。為了實(shí)現(xiàn)這一點(diǎn)，首先，讓我們編寫(xiě)一個(gè)發(fā)送電子郵件的函數(shù)。我們將使用Gmail發(fā)送電子郵件。
以下是發(fā)送電子郵件的函數(shù)：
def send_email(
body,
from_email=FROM_EMAIL,
to_emails=RECIPENT_EMAIL,
subject=subject,
api=EMAIL_API_KEY,
):
msg = MIMEMultipart()
msg["From"] = from_email
msg["To"] = to_emails
msg["Subject"] = subject
msg.attach(MIMEText(body, "html"))
try:
smtp_server = smtplib.SMTP("smtp.gmail.com", 587)
smtp_server.starttls()
smtp_server.login(from_email, api)
smtp_server.sendmail(from_email, to_emails, msg.as_string())
smtp_server.quit()
print("Email sent.")
except Exception as e:
print("Email not sent", e)但請(qǐng)注意，我們需要將諸如EMAIL API KEY、SENDER EMAIL等秘密憑證存儲(chǔ)到一個(gè)環(huán)境文件中?；诖?，請(qǐng)?jiān)谀愕捻?xiàng)目根目錄中創(chuàng)建一個(gè)新文件（命名為.env）。在.env文件中，添加以下示例內(nèi)容。
EMAIL_API_KEY=your Gmail app password goes here
EMAIL_ACCOUNT=your Gmail account which you created app password goes here
RECIPENT_EMAIL=the email address you will be sending the report email goes here.現(xiàn)在，讓我們繼續(xù)實(shí)現(xiàn)回調(diào)函數(shù)（on_train_end），該函數(shù)將在模型訓(xùn)練成功完成后觸發(fā)發(fā)送電子郵件功能。
def on_train_end(trainer):
trainer_epoch = trainer.epoch
trainer_metrics = trainer.metrics
current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
end_time = datetime.now()
time_taken = end_time - start_time
hours, remainder = divmod(time_taken.total_seconds(), 3600)
minutes, seconds = divmod(remainder, 60)
time_taken_str = ""
if int(hours) > 0:
time_taken_str += f"{int(hours)} hr "
if int(minutes) > 0:
time_taken_str += f"{int(minutes)} mins "
if int(seconds) > 0:
time_taken_str += f"{int(seconds)} secs"
time_taken_str = time_taken_str.strip()
body = f"""
<html>
<head>
<style>
table, th, td {{
border: 1px solid black;
border-collapse: collapse;
padding: 5px;
}}
</style>
</head>
<body>
<h1>Training Report</h1>
<p>Date and Time: {current_time}</p>
<p>Total Epoch Trained: {trainer_epoch + 1} </p>
<p>Time Taken to Train Model: {time_taken_str} </p>
<table>
<tr>
<th>Metric</th>
<th>Value</th>
</tr>
{''.join([f'<tr><td>{k}</td><td>{v:.2f}</td></tr>' for k, v in trainer_metrics.items()])}
</table>
</body>
</html>
"""
send_email(body)以上回調(diào)函數(shù)將在模型訓(xùn)練完成后向指定收件人發(fā)送報(bào)告郵件?，F(xiàn)在，我們已經(jīng)編寫(xiě)了所有必要的函數(shù)，將它們?nèi)糠庋b在一個(gè)名為ModelTraining的類中是一個(gè)好主意。所以，我們training.py文件中的完整代碼現(xiàn)在應(yīng)該如下所示：
import os
from datetime import datetime
from dotenv import find_dotenv, load_dotenv
import torch
from ultralytics import YOLO
import smtplib
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
load_dotenv(find_dotenv())
EMAIL_API_KEY = os.getenv("EMAIL_API_KEY")
FROM_EMAIL = os.getenv("EMAIL_ACCOUNT")
RECIPIENT_EMAIL = os.getenv("RECIPIENT_EMAIL")
subject = "Model Training Completed"
class ModelTraining:
def __init__(self):
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.start_time = None
self.end_time = None
def send_email(
self,
body,
from_email=FROM_EMAIL,
to_emails=RECIPIENT_EMAIL,
subject=subject,
api=EMAIL_API_KEY,
):
msg = MIMEMultipart()
msg["From"] = from_email
msg["To"] = to_emails
msg["Subject"] = subject
msg.attach(MIMEText(body, "html"))
try:
smtp_server = smtplib.SMTP("smtp.gmail.com", 587)
smtp_server.starttls()
smtp_server.login(from_email, api)
smtp_server.sendmail(from_email, to_emails, msg.as_string())
smtp_server.quit()
print("Email sent.")
except Exception as e:
print("Email not sent", e)
def on_train_start(self, trainer):
self.start_time = datetime.now()
def on_train_epoch_end(self, trainer):
curr_epoch = trainer.epoch + 1
text = f"Epoch Number: {curr_epoch}/{trainer.epochs} finished"
print(text)
print("-" * 50)
def on_train_end(self, trainer):
trainer_epoch = trainer.epoch
trainer_metrics = trainer.metrics
current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
self.end_time = datetime.now()
time_taken = self.end_time - self.start_time
hours, remainder = divmod(time_taken.total_seconds(), 3600)
minutes, seconds = divmod(remainder, 60)
time_taken_str = ""
if int(hours) > 0:
time_taken_str += f"{int(hours)} hr "
if int(minutes) > 0:
time_taken_str += f"{int(minutes)} mins "
if int(seconds) > 0:
time_taken_str += f"{int(seconds)} secs"
time_taken_str = time_taken_str.strip()
body = f"""
<html>
<head>
<style>
table, th, td {{
border: 1px solid black;
border-collapse: collapse;
padding: 5px;
}}
</style>
</head>
<body>
<h1>Training Report</h1>
<p>Date and Time: {current_time}</p>
<p>Total Epochs Trained: {trainer_epoch + 1} </p>
<p>Time Taken to Train Model: {time_taken_str} </p>
<table>
<tr>
<th>Metric</th>
<th>Value</th>
</tr>
{''.join([f'<tr><td>{k}</td><td>{v:.2f}</td></tr>' for k, v in trainer_metrics.items()])}
</table>
</body>
</html>
"""
self.send_email(body)
def train_yolov8_model(self, config_path, num_epochs, training_result_dir):
model = YOLO("yolov8x.pt")
model.add_callback("on_train_start", self.on_train_start)
model.add_callback("on_train_epoch_end", self.on_train_epoch_end)
model.add_callback("on_train_end", self.on_train_end)
model.start_time = datetime.now()
start_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
# Train the model
model.train(
data=config_path,
name="Yolo_Model_Training",
project=training_result_dir,
task="detect",
epochs=num_epochs,
patience=20,
batch=16,
cache=True,
imgsz=640,
iou=0.5,
augment=True,
degrees=25.0,
fliplr=0.0,
lr0=0.0001,
optimizer="Adam",
device=self.device,
)
model.end_time = datetime.now()
if __name__ == "__main__":
model_training = ModelTraining()
# Load the dataset configuration file
current_dir = os.path.dirname(os.path.abspath(__file__))
config_path = os.path.join(current_dir, "data.yaml")
num_epochs = 40 # Change it to any number of epochs you want.
training_result_path = "./results"
os.makedirs(training_result_path, exist_ok=True)
model_training.train_yolov8_model(config_path, num_epochs, training_result_path)完整的項(xiàng)目結(jié)構(gòu)應(yīng)該如下所示：
yolo_with_callback/
│
├── dataset/ # Directory containing dataset files
│
├── env/ # python virtual environment directory
│
│── .env # Environment variables file containing secret keys
├── results/ # Directory for storing training results
│
├── data.yaml # Dataset configuration file
│
├── requirements.txt # File listing required Python packages
│
└── training.py # Main script for model training現(xiàn)在，你已經(jīng)完成了實(shí)現(xiàn)，可以繼續(xù)運(yùn)行training.py代碼。訓(xùn)練完成后，訓(xùn)練結(jié)果報(bào)告將發(fā)送到指定的收件人郵箱。
