yolov5基础
qinsz 2/17/2024 torchyolov5
# 一、环境
# 1、CUDA环境（可选）
# 1.1 CUDA Toolkit
- CMD中通过
nvidia-smi查看CUDA版本 - 前往nvidia官网下载并安装对应版本的CUDA Toolkit:
- https://developer.nvidia.com/cuda-downloads
- 默认安装会自动添加nvcc环境变量，使用
nvcc --version可查看到CUA版本- nvcc.exe位于
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8\bin
- nvcc.exe位于
# 2. 安装YOLO11
pip install ultralytics onnxslim onnxruntime onnx
1
- 其中
onnxslim onnxruntime onnx是需要导出onnx模型时需要用到的，可根据需要决定是否进行安装
# 3. 安装CUDA版torch（可选）
希望使用GPU需要安装兼容CUDA版本的torch库（前提计算机中的CUDA环境已经正常）
卸载已有的torch库
pip uninstall torch torchvision torchaudio1前往torch官网获取CUDA版torch库的信息:
https://pytorch.org/get-started/locally/- 在页面中选择正确的环境获得库信息然后安装:
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu1261- Tips:
torch-2.6.0+cu126-cp310-cp310-win_amd64.whl很大，下载速度慢，可以用迅雷单独下载，然后执行pip的本地安装:
pip install /xxxpath/torch-2.6.0+cu126-cp310-cp310-win_amd64.whl1 - Tips:
- 在页面中选择正确的环境获得库信息然后安装:
# 二、训练模型
# 1. 准备数据
- 训练数据图片：/path/.../my-datasets/images/train/
- 训练数据label: /path/.../my-datasets/labels/train/
- 验证数据图片：/path/.../my-datasets/images/val/
- 验证数据label: /path/.../my-datasets/labels/val/
# 2. 数据集配置文件
# /.../project/fire_detect_data_set.yaml
path: /path/.../my-datasets
train: images/train # 指定训练集image位置，会自动在上级路径中找labels文件夹中的同名label
val: images/val # 制定测试集image位置
names:
0: fire # 0-目标标记，fire-目标名称
nc: 1 # 类别数量
1
2
3
4
5
6
7
8
2
3
4
5
6
7
8
# 3. 开始训练
# /.../project/fire_detect_train.py
from ultralytics import YOLO
if __name__ == '__main__':
model = YOLO('yolo11n.pt') # 选择YOLO内置的一个模型结构
train_results = model.train(
data='fire_detect_data_set.yaml', # 指定训练数据的配置文件
epochs=200,
batch=32,
device='cuda'
)
1
2
3
4
5
6
7
8
9
10
11
12
2
3
4
5
6
7
8
9
10
11
12
训练完成后自动将best.pt和last.pt文件保存到/.../project/runs/detect/trainx/weights/