37 lines
1.5 KiB
Python
37 lines
1.5 KiB
Python
import argparse
|
|
from ultralytics import YOLO
|
|
|
|
parser = argparse.ArgumentParser(description="YOLO26 岩性分类训练")
|
|
parser.add_argument("--data", default="custom_dataset", help="数据集路径")
|
|
parser.add_argument("--model", default="yolo26n-cls.pt", help="预训练权重路径")
|
|
parser.add_argument("--epochs", type=int, default=40, help="训练轮数")
|
|
parser.add_argument("--imgsz", type=int, default=224, help="输入图片尺寸")
|
|
parser.add_argument("--batch", type=int, default=32, help="批次大小")
|
|
parser.add_argument("--lr0", type=float, default=0.001, help="初始学习率")
|
|
parser.add_argument("--device", default="cuda", help="训练设备")
|
|
parser.add_argument("--patience", type=int, default=10, help="早停耐心值")
|
|
parser.add_argument("--workers", type=int, default=0, help="数据加载线程数")
|
|
parser.add_argument("--cos-lr", action="store_true", default=True, help="使用余弦学习率")
|
|
parser.add_argument("--label-smoothing", type=float, default=0.15, help="标签平滑")
|
|
parser.add_argument("--plots", action="store_true", help="绘制训练曲线")
|
|
args = parser.parse_args()
|
|
|
|
model = YOLO(args.model)
|
|
|
|
model.train(
|
|
data=args.data,
|
|
epochs=args.epochs,
|
|
imgsz=args.imgsz,
|
|
batch=args.batch,
|
|
lr0=args.lr0,
|
|
patience=args.patience,
|
|
cos_lr=args.cos_lr,
|
|
augment=True,
|
|
label_smoothing=args.label_smoothing,
|
|
workers=args.workers,
|
|
task="classify",
|
|
device=args.device,
|
|
plots=args.plots,
|
|
exist_ok=True,
|
|
)
|