在Ubuntu系统上使用YOLOv5训练自己的模型,你需要遵循以下步骤。这里我将详细说明如何从准备数据集到训练模型的整个过程。
步骤 1: 安装依赖项
首先,确保你的Ubuntu系统上安装了Python、PyTorch和必要的库。你可以使用以下命令安装这些依赖项:
sudo apt update
sudo apt install python3-pip
pip3 install torch torchvision torchaudio
pip3 install numpy matplotlib
pip3 install opencv-python
pip3 install Pillow
pip3 install scipy
步骤 2: 克隆YOLOv5仓库
接下来,克隆YOLOv5的GitHub仓库到你的系统中:
git clone https://github.com/ultralytics/yolov5.git
cd yolov5
步骤 3: 准备数据集
你需要准备自己的数据集。数据集应该包含图片和相应的标注文件(通常是YOLO格式的txt文件)。例如,你的文件夹结构可能是这样的:
/dataset
/images
image1.jpg
image2.jpg
…
/labels
image1.txt
image2.txt
…
每个.txt文件包含与图像对应的标注信息,
安装labelimg来进行图片信息标注,下载地址
pip install labelimg -i https://pypi.tuna.tsinghua.edu.cn/simple
输出结果例如:
0 0.1 0.2 0.3 0.4 # class x_center y_center width height
1 0.6 0.7 0.1 0.2 # class x_center y_center width height
步骤 4: 创建数据集配置文件
在yolov5目录中,编辑data文件夹下的hyp.yaml文件来配置你的数据集路径和类别:
train: …/dataset/images/train/ # 训练图像路径
val: …/dataset/images/val/ # 验证图像路径
nc: 2 # 类别数,例如2个类别(人、车)
names: [‘person’, ‘car’] # 类别名称列表
步骤 5: 训练模型
使用以下命令开始训练你的模型:
python train.py --img 640 --batch 16 --epochs 50 --data data/hyp.yaml --weights yolov5s.pt --name my_experiment --cache
这里是一些参数的解释:
–img:输入图像的大小。
–batch:批量大小。
–epochs:训练周期数。
–data:数据集配置文件路径。
–weights:预训练权重,例如yolov5s.pt。如果你从头开始训练,可以使用yolov5s.pt。对于从头开始,你也可以使用yolov5s.yaml配置文件来指定网络结构。
–name:实验名称,用于保存模型和结果。
–cache:启用缓存以提高性能。
步骤 6: 验证和测试模型
训练完成后,你可以使用以下命令来评估你的模型:
python val.py --weights runs/train/my_experiment/weights/best.pt --data data/hyp.yaml --img 640 --conf 0.4 --iou 0.5 --device 0 --half # 使用最佳权重进行评估
步骤 7: 进行推理和测试(可选)
最后,你可以使用以下命令来测试你的模型:
python detect.py --weights runs/train/my_experiment/weights/best.pt --source data/dataset/images/test --img 640 --conf 0.4 --iou 0.5 --device 0 --half # 对测试图像进行检测
这样,你就完成了在Ubuntu上使用YOLOv5训练自己的模型的过程。确保根据你的具体需求调整参数和配置。