File size: 499 Bytes
7aad26f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
from ultralytics import YOLO
if __name__ == "__main__":
# Load YOLO model
model = YOLO("yolov8s.yaml")
# Train the model
results = model.train(
data="datasets/data.yaml",
epochs=100,
imgsz=640,
workers=0,
batch=16,
device="cuda",
plots = False
)
# Save model
model_path = "runs/detect/train/weights/best.pt"
model.export(format="torchscript")
print(f" Model saved at: {model_path}")
|