from ultralytics import YOLO

def main():
    # 加载预训练模型
    model = YOLO("yolov8n.pt")  # 使用YOLOv8 nano版本

    # 开始训练
    results = model.train(
        data="datasets/猫狗/data.yaml",
        epochs=100,
        batch=8,
        imgsz=640,
        device="cpu"  # 使用GPU，如果是CPU则设为"cpu"
    )

if __name__ == '__main__':
    main()