yolov11 训练自己的模型

1.conda create -n testgpu python=3.8

2.

conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.7 -c pytorch -c nvidia

3.

pip install ultralytics -i Simple Index

4.

pip install numpy==1.26.4 -i Simple Index

训练代码 device=1就是gpu

from ultralytics import YOLO
import os
import torch
if __name__ == '__main__':
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")
    model = YOLO(".\\yolo11n.pt")
    current_dir = os.path.dirname(os.path.abspath(__file__))
    # 构建数据集配置文件的相对路径
    print(current_dir)
    data_path = os.path.join(current_dir, "appgpu.yaml")
    result = model.train(
        data=data_path, 
        epochs=100, 
        imgsz=640,