CV_AG/Train_model.py

19 lines
804 B
Python

from ultralytics import YOLO
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
os.environ["CUDA_VISIBLE_DEVICES"]="0"
if __name__ == '__main__':
model = YOLO('yolo11n.pt')
results = model.train(
data = "D:\AIM\lemon\AIM.v3i.yolov11\data.yaml", # Path to the dataset (YAML file for COCO128)
epochs = 50, # Number of epochs for fine-tuning
imgsz = 640, # Image size for training
batch = 16, # Batch size
lr0 = 0.005, # Initial learning rate
workers = 4, # Number of data loading workers
device = 0 # Specify the device (0 = first GPU, 'cpu' for CPU)
)