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()
|