from ultralytics import YOLO

def main():
    model = YOLO("models/unified-nano.pt")
    model.train(
        data="./datasets/pieces/unified/data.yaml",
        epochs=150,
        patience=20,
        imgsz=640,
        batch=18,
        save_period=10,
        project="result",
        name="unified-nano-refined",
        exist_ok=True,
        device = 0
    )

if __name__ == "__main__":
    main()