Zuxin Dai
34 lines
1.3 KiB
Python
34 lines
1.3 KiB
Python
from ultralytics import YOLO
|
|
import os
|
|
|
|
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
|
|
|
|
def main():
|
|
# 初始化 Weights & Biases
|
|
import wandb
|
|
wandb.login(key='7cfbcf76a18a8441b04eb5d7adb988e69a79705e') # 替换为你的 API 密钥
|
|
wandb.init(project="YOLO-Training", name="YOLOv11_finetune", mode="online")
|
|
|
|
# 加载 YOLO 模型
|
|
model = YOLO('yolo11n.pt')
|
|
|
|
# 开始训练
|
|
model.train(
|
|
data=r"D:\AIM\lemon\lemon_quality_dataset_YOLO11\data.yaml", # 数据集配置文件路径
|
|
epochs=2, # 训练轮次
|
|
imgsz=640, # 图像大小
|
|
batch=16, # 批量大小
|
|
lr0=0.005, # 初始学习率
|
|
workers=0, # 设置为 0，禁用多进程数据加载
|
|
device=0, # 设备 (0 = 第一块 GPU, 'cpu' = CPU)
|
|
project="YOLO-Training", # W&B 项目名称
|
|
name="YOLOv11_finetune", # W&B 实验名称
|
|
exist_ok=True # 如果目录存在是否覆盖
|
|
)
|
|
|
|
if __name__ == '__main__':
|
|
import torch.multiprocessing as mp
|
|
mp.freeze_support() # 解决 Windows 下的多进程启动问题
|
|
main()
|