26 lines
1.3 KiB
Python
26 lines
1.3 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.load('yolov8n.pt') # loading pretrain weights
|
||
model.train(data='G:/skin-cancer-detection',
|
||
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',
|
||
) |