27 lines
1.4 KiB
Python
27 lines
1.4 KiB
Python
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__':
|
||
# 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',
|
||
cache=False,
|
||
imgsz=224,
|
||
epochs=2,
|
||
batch=64,
|
||
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',
|
||
) |