-
Notifications
You must be signed in to change notification settings - Fork 436
/
Copy pathtrain_yolo11.py
33 lines (26 loc) · 1.04 KB
/
train_yolo11.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
import sys
import argparse
import os
from ultralytics import YOLO
def main(opt):
yaml = opt.cfg
model = YOLO(yaml) # 直接加载yaml文件训练
# model = YOLO(weights) # 直接加载权重文件进行训练
# model = YOLO(yaml).load(weights) # 加载yaml配置文件的同时,加载权重进行训练
model.info()
results = model.train(data='coco128.yaml', # 训练参数均可以重新设置
epochs=2,
imgsz=640,
workers=2,
batch=2,
# ...在这里添加需要修改的参数
)
def parse_opt(known=False):
parser = argparse.ArgumentParser()
parser.add_argument('--cfg', type=str, default='ultralytics/cfg/models/11/yolo11.yaml', help='initial weights path')
parser.add_argument('--weights', type=str, default='', help='')
opt = parser.parse_known_args()[0] if known else parser.parse_args()
return opt
if __name__ == "__main__":
opt = parse_opt()
main(opt)