1. 环境安装

    1)  训练环境基于:

        PyTorch  1.10.0        Python  3.8(ubuntu20.04)               CUDA  11.3

   2) 编译依赖包安装

       pip install -r requirements.txt  , requirements.txt  文件如下

# Ultralytics requirements
# Example: pip install -r requirements.txt

# Base ----------------------------------------
matplotlib>=3.3.0
numpy>=1.22.2 # pinned by Snyk to avoid a vulnerability
opencv-python>=4.6.0
pillow>=7.1.2
pyyaml>=5.3.1
requests>=2.23.0
scipy>=1.4.1
torch>=1.8.0
torchvision>=0.9.0
tqdm>=4.64.0
polars>=1.8.2
# Logging -------------------------------------
# tensorboard>=2.13.0
# dvclive>=2.12.0
# clearml
# comet

# Plotting ------------------------------------
pandas>=1.1.4
seaborn>=0.11.0

# Export --------------------------------------
# coremltools>=7.0  # CoreML export
# onnx>=1.12.0  # ONNX export
# onnxsim>=0.4.1  # ONNX simplifier
# nvidia-pyindex  # TensorRT export
# nvidia-tensorrt  # TensorRT export
# scikit-learn==0.19.2  # CoreML quantization
# tensorflow>=2.4.1  # TF exports (-cpu, -aarch64, -macos)
# tflite-support
# tensorflowjs>=3.9.0  # TF.js export
# openvino-dev>=2023.0  # OpenVINO export

# Extras --------------------------------------
psutil  # system utilization
py-cpuinfo  # display CPU info
thop>=0.1.1  # FLOPs computation
# ipython  # interactive notebook
# albumentations>=1.0.3  # training augmentations
# pycocotools>=2.0.6  # COCO mAP
# roboflow

2. 训练文件编写

 ultralytics-yolo26/user_train.py,指定自己数据集的data.yaml文件

from ultralytics import YOLO

# Load a pretrained YOLO26n model
model = YOLO("yolo26n.yaml")
# model = YOLO("yolo26n.pt")  #从预训练模型训练
# Train the model on the strawberry dataset for 100 epochs
train_results = model.train(
    data="/root/strawberry/data.yaml",  # Path to dataset configuration file
    epochs=100,  # Number of training epochs
    imgsz=640  # Image size for training
)    #    resume=True #断点续训

执行命令训练: python  user_train.py,自主训练草莓数据集训练结束结果如下

3. 基于训练好的模型进行预测

   ultralytics-yolo26/user_train.py,user_train.py文件如下

from ultralytics import YOLO

# Load a pretrained YOLOv8n model
model = YOLO('runs/detect/train2/weights/best.pt') #已经训练好的模型路径

# Define path to the image file
source = '/root/ultralytics-yolo26/predict_user_image' #待预测的数据保存路径
#source = 0
# Run inference on the source
results = model(source, mode='predict', save=True)  # list of Results objects


执行命令训练: python  user_predict.py

Logo

DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。

更多推荐