19 lines
804 B
Python
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)
|
|
) |