ultralytics/train.py

26 lines
1.3 KiB
Python
Raw Normal View History

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-02-25 07:10:10 +00:00
model = YOLO('ultralytics/cfg/models/v8/yolov8n-cls.yaml')
2025-02-25 03:58:34 +00:00
# model.load('yolov8n.pt') # loading pretrain weights
2025-02-25 07:10:10 +00:00
model.train(data='G:/skin-cancer-detection',
2025-02-25 03:58:34 +00:00
cache=False,
imgsz=640,
epochs=300,
batch=32,
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',
)