2025-02-25 03:58:34 +00:00
|
|
|
|
import warnings, os
|
|
|
|
|
|
# os.environ["CUDA_VISIBLE_DEVICES"]="-1" # 代表用cpu训练 不推荐！没意义！ 而且有些模块不能在cpu上跑
|
|
|
|
|
|
# os.environ["CUDA_VISIBLE_DEVICES"]="0" # 代表用第一张卡进行训练 0：第一张卡 1：第二张卡
|
|
|
|
|
|
# 多卡训练参考<YOLOV8V10配置文件.md>下方常见错误和解决方案
|
|
|
|
|
|
warnings.filterwarnings('ignore')
|
|
|
|
|
|
from ultralytics import YOLO
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
2025-03-01 10:30:01 +00:00
|
|
|
|
# model = YOLO('ultralytics/cfg/models/v8/yolov8n-cls.yaml') # 单任务学习
|
|
|
|
|
|
# model.train(data='G:/dataset/split',
|
|
|
|
|
|
model = YOLO('ultralytics/cfg/models/v8/yolov8-mtlcls.yaml',task='MTL') # 多任务学习
|
|
|
|
|
|
model.train(data='G:/dataset/test/ml.yaml',
|
2025-02-25 03:58:34 +00:00
|
|
|
|
cache=False,
|
2025-03-01 10:30:01 +00:00
|
|
|
|
imgsz=224,
|
|
|
|
|
|
epochs=2,
|
|
|
|
|
|
batch=64,
|
2025-02-25 03:58:34 +00:00
|
|
|
|
close_mosaic=0,
|
|
|
|
|
|
workers=8, # Windows下出现莫名其妙卡主的情况可以尝试把workers设置为0
|
|
|
|
|
|
optimizer='SGD', # using SGD
|
|
|
|
|
|
# device='0,1', # 指定显卡和多卡训练参考<YOLOV8V10配置文件.md>下方常见错误和解决方案
|
|
|
|
|
|
# patience=0, # set 0 to close earlystop.
|
|
|
|
|
|
# resume=True, # 断点续训,YOLO初始化时选择last.pt,例如YOLO('last.pt')
|
|
|
|
|
|
# amp=False, # close amp
|
|
|
|
|
|
# fraction=0.2,
|
|
|
|
|
|
project='runs/train',
|
|
|
|
|
|
name='exp',
|
|
|
|
|
|
)
|
