26 lines
1022 B
Python
26 lines
1022 B
Python
from ultralytics import YOLO
|
|
import torch
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
# # 单纯安全帽数据集训练
|
|
# model = YOLO('yolov8n.yaml')
|
|
#
|
|
# results = model.train(data='./5000_HeadHelmet.yaml', epochs=100, batch=1, patience=20, imgsz=640, save=True,
|
|
# workers=0, device=0, project='5000_HeadHelmet', name='train', verbose=True)
|
|
|
|
# model = YOLO('yolov8n_hhfs.yaml')
|
|
|
|
# results = model.train(data='newdataset.yaml', epochs=100, batch=4, patience=20, imgsz=640, save=True,
|
|
# workers=0, device=0, project='10000_newdataset', name='train', verbose=True)
|
|
|
|
# print(torch.cuda.is_available())
|
|
|
|
model = YOLO('./yolov8n.yaml')
|
|
|
|
results = model.train(data='/home/admin-root/haotian/python哈汽锻8安全帽识别/dataset_20250428_BAC.yaml', epochs=200, batch=32, patience=20, imgsz=640, save=True,
|
|
workers=0, device=0, project='HelmetHeadBAC', name='train', verbose=True)
|
|
|
|
# print(model) |