侧边栏壁纸
博主头像
天涯小筑

行动起来,活在当下

  • 累计撰写 4 篇文章
  • 累计创建 2 个标签
  • 累计收到 0 条评论

目 录CONTENT

文章目录

ubuntu下快速搭建docker环境训练yolov5数据集

# 参考文档

yolov5-github

yolov5-github-训练文档

csdn训练博客

一、配置环境

1.1 更新国内镜像源

可以使用如下工具更新镜像源:

sudo wget http://fishros.com/install -O fishros && . fishros

1.2 安装依赖包

  • 安装必要的环境依赖包

sudo apt-get install -y build-essential ubuntu-drivers-common net-tools python3 python-is-python3 python3-pip
  • 修改pip源为清华源

pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

或者修改配置文件替换pipe国内源地址

mkdir -p ~/.pip && tee ~/.pip/pip.conf > /dev/null <<EOF
[global]
index-url = https://pypi.tuna.tsinghua.edu.cn/simple
[install]
trusted-host = pypi.tuna.tsinghua.edu.cn
EOF

执行验证:

pip config list
# global.index-url='https://pypi.tuna.tsinghua.edu.cn/simple'
# install.trusted-host='https://pypi.tuna.tsinghua.edu.cn'

1.3 安装docker

具体安装步骤参考ubuntu安装docker官方文档

简单直接的安装方式参考如下:

 sudo apt install -y docker.io
 # 将当前用户加入docker组
 sudo usermod -aG docker ${USER}

1.4 拉取pytorch docker镜像

前往pytorch 官方docker镜像寻找自己合适版本,yolov5要求1.8以上版本,我拉取1.13版本,执行命令:

sudo docker pull pytorch/pytorch:1.13.1-cuda11.6-cudnn8-runtime

1.4 安装nvidia驱动

桌面版参考链接

服务器版参考链接

我们使用pytorch-docker环境无需安装cudaNVIDIA驱动简单安装如下

  1. 禁用nouveau驱动

编辑 /etc/modprobe.d/blacklist-nouveau.conf 文件,添加以下内容:

blacklist nouveau
blacklist lbm-nouveau
options nouveau modeset=0
alias nouveau off
alias lbm-nouveau off
  1. 关闭nouveau

   echo options nouveau modeset=0 | sudo tee -a /etc/modprobe.d/nouveau-kms.conf
  1. 重新生成内核并重启

sudo update-initramfs -u
sudo reboot
  1. 重启后验证

重启后,执行lsmod | grep nouveau如果没有屏幕输出,说明禁用nouveau成功

  1. 查找推荐驱动

ubuntu-drivers devices
# 输出如下
# modalias : pci:v000010DEd00001EB8sv000010DEsd000012A2bc03sc02i00
# vendor   : NVIDIA Corporation
# model    : TU104GL [Tesla T4]
# driver   : nvidia-driver-450-server - distro non-free
# driver   : nvidia-driver-525-server - distro non-free
# driver   : nvidia-driver-535-server - distro non-free
# driver   : nvidia-driver-418-server - distro non-free
# driver   : nvidia-driver-525 - distro non-free
# driver   : nvidia-driver-470 - distro non-free
# driver   : nvidia-driver-470-server - distro non-free
# driver   : nvidia-driver-535 - distro non-free recommended
# driver   : xserver-xorg-video-nouveau - distro free builtin
  1. 安装推荐的驱动程序

根据自己系统选择安装,安装完成后重启

   sudo apt install nvidia-driver-535-server
  1. 重启后验证

nvidia-smi 命令能够输出显卡信息则验证成功

二、训练数据集

2.1 下载yolov5代码

前往github下载代码,或者准备自己的yolov5训练代码,如果是拷贝他人代码,将git目录删除,否则后续训练时检查git信息会报错。

git clone git@github.com:ultralytics/yolov5.git

2.2 启动进入pytorch-docker

# 映射宿主机地址到docker内部,根据显卡实际情况指定显存容量
docker run -v /home/lishi/object-detect/:/workspace --gpus all --ipc=host -p 6006:6006 -it pytorch/pytorch:1.13.1-cuda11.6-cudnn8-runtime /bin/bash

后续都将在docker中执行;

2.3 安装依赖项

docker下进入yolov5代码目录下将request.txtopencv注释掉然后执行依赖项安装

pip3 install -r requirements.txt -i http://mirrors.aliyun.com/pypi/simple/ --trusted-host mirrors.aliyun.com

继续安装opencv-python-headless版本opencv

pip3 install opencv-python-headless -i http://mirrors.aliyun.com/pypi/simple/ --trusted-host mirrors.aliyun.com

2.4 准备训练数据集

2.4.1 VOC数据集

VOC数据集需要转换为YOLO数据集训练,转换下方转换代码,修改CLASSES和PATH

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import random
from shutil import copyfile
from PIL import Image

# 只要改下面的CLASSES和PATH就可以了,其他的不用改,这个脚本会自动划分数据集,生成YOLO格式的标签文件
# 分类名称  这里改成数据集的分类名称,一定要改!!!请查看数据集目录下的txt文件
CLASSES = ["belt", "nobelt"]
# 数据集目录 这里改成数据集的根目录,根目录下有两个文件夹Annotations和JPEGImages,一定要改!!!
PATH = r'/home/lishi/object-detect/Helmet/data-images'
# 训练集占比80% 训练集:验证集=8:2 这里划分数据集 不用改
TRAIN_RATIO = 80
def clear_hidden_files(path):
    dir_list = os.listdir(path)
    for i in dir_list:
        abspath = os.path.join(os.path.abspath(path), i)
        if os.path.isfile(abspath):
            if i.startswith("._"):
                os.remove(abspath)
        else:
            clear_hidden_files(abspath)


def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)


def convert_annotation(image_id):
    # Assuming the image format is jpg
    image_path = os.path.join(image_dir, f"{image_id}.jpg")
    img = Image.open(image_path)
    w, h = img.size
    in_file = open(PATH+'/Annotations/%s.xml' % image_id, encoding='utf-8')
    out_file = open(PATH+'/YOLOLabels/%s.txt' %
                    image_id, 'w', encoding='utf-8')
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    # w = int(size.find('width').text)
    # h = int(size.find('height').text)
    difficult = 0
    for obj in root.iter('object'):
        if obj.find('difficult'):
            difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in CLASSES or int(difficult) == 1:
            continue
        cls_id = CLASSES.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text))
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " +
                       " ".join([str(a) for a in bb]) + '\n')
    in_file.close()
    out_file.close()


wd = os.getcwd()
wd = os.getcwd()

work_sapce_dir = os.path.join(wd, PATH+"/")

annotation_dir = os.path.join(work_sapce_dir, "Annotations/")
if not os.path.isdir(annotation_dir):
    os.mkdir(annotation_dir)
clear_hidden_files(annotation_dir)
image_dir = os.path.join(work_sapce_dir, "JPEGImages/")
if not os.path.isdir(image_dir):
    os.mkdir(image_dir)
clear_hidden_files(image_dir)
yolo_labels_dir = os.path.join(work_sapce_dir, "YOLOLabels/")
if not os.path.isdir(yolo_labels_dir):
    os.mkdir(yolo_labels_dir)
clear_hidden_files(yolo_labels_dir)

yolov5_train_dir = os.path.join(work_sapce_dir, "train/")
if not os.path.isdir(yolov5_train_dir):
    os.mkdir(yolov5_train_dir)
clear_hidden_files(yolov5_train_dir)
yolov5_images_train_dir = os.path.join(yolov5_train_dir, "images/")
if not os.path.isdir(yolov5_images_train_dir):
    os.mkdir(yolov5_images_train_dir)
clear_hidden_files(yolov5_images_train_dir)
yolov5_labels_train_dir = os.path.join(yolov5_train_dir, "labels/")
if not os.path.isdir(yolov5_labels_train_dir):
    os.mkdir(yolov5_labels_train_dir)
clear_hidden_files(yolov5_labels_train_dir)

yolov5_test_dir = os.path.join(work_sapce_dir, "val/")
if not os.path.isdir(yolov5_test_dir):
    os.mkdir(yolov5_test_dir)
clear_hidden_files(yolov5_test_dir)
yolov5_images_test_dir = os.path.join(yolov5_test_dir, "images/")
if not os.path.isdir(yolov5_images_test_dir):
    os.mkdir(yolov5_images_test_dir)
clear_hidden_files(yolov5_images_test_dir)
yolov5_labels_test_dir = os.path.join(yolov5_test_dir, "labels/")
if not os.path.isdir(yolov5_labels_test_dir):
    os.mkdir(yolov5_labels_test_dir)
clear_hidden_files(yolov5_labels_test_dir)


train_file = open(os.path.join(wd, "yolov5_train.txt"), 'w', encoding='utf-8')
test_file = open(os.path.join(wd, "yolov5_valid.txt"), 'w', encoding='utf-8')
train_file.close()
test_file.close()
train_file = open(os.path.join(wd, "yolov5_train.txt"), 'a', encoding='utf-8')
test_file = open(os.path.join(wd, "yolov5_valid.txt"), 'a', encoding='utf-8')
list_imgs = os.listdir(image_dir)  # list image files
prob = random.randint(1, 100)
print("数据集: %d个" % len(list_imgs))
for i in range(0, len(list_imgs)):
    path = os.path.join(image_dir, list_imgs[i])
    if os.path.isfile(path):
        image_path = image_dir + list_imgs[i]
        voc_path = list_imgs[i]
        (nameWithoutExtention, extention) = os.path.splitext(
            os.path.basename(image_path))
        (voc_nameWithoutExtention, voc_extention) = os.path.splitext(
            os.path.basename(voc_path))
        annotation_name = nameWithoutExtention + '.xml'
        annotation_path = os.path.join(annotation_dir, annotation_name)
        label_name = nameWithoutExtention + '.txt'
        label_path = os.path.join(yolo_labels_dir, label_name)
    prob = random.randint(1, 100)
    print("Probability: %d" % prob, i, list_imgs[i])
    if (prob < TRAIN_RATIO):
        # train dataset
        if os.path.exists(annotation_path):
            train_file.write(image_path + '\n')
            convert_annotation(nameWithoutExtention)  # convert label
            copyfile(image_path, yolov5_images_train_dir + voc_path)
            copyfile(label_path, yolov5_labels_train_dir + label_name)
    else:
        # test dataset
        if os.path.exists(annotation_path):
            test_file.write(image_path + '\n')
            convert_annotation(nameWithoutExtention)  # convert label
            copyfile(image_path, yolov5_images_test_dir + voc_path)
            copyfile(label_path, yolov5_labels_test_dir + label_name)
train_file.close()
test_file.close()

2.4.2 YOLO格式数据集

yolo格式数据集通常为如下结构:

这种数据集不用做任何处理,只需将dataset.yaml内容修改为pytorch-docker环境下的绝对路径即可,如下参考:

train: /workspace/Helmet/image-data/train/images
val: /workspace/Helmet/image-data/valid/images
nc: 2
names:
  - belt
  - nobelt

train: 容器环境下训练数据集绝对路径;

val: 容器环境下验证数据集绝对路径;

nc: 数据集类别数量

names: 数据集类别名称

修改后将dataset.yaml拷贝到源码data目录下,后续在train.py训练代码中--data参数设置为data/dataset.yaml

2.5 修改模型文件

models下有5个模型,smlx需要训练的时间依次增加,按照需求选择一个文件进行修改即可,我选择yolov5s.yaml

只需将nc改为实际值即可;

yolov5s.yaml修改nc为实际值;

2.6 修改训练tran.py

这里需要对train.py文件内的参数进行修改,weightscfgdata按照自己所需文件的路径修改,weights如果使用参考博客的文件,将yolov5s.pt下载放到代码根目录下即可,如果使用官方则无需修改,会自行下载。具体参数含义,查看官方文档。我修改内容如下:

主要修改了以下几项:

cfg :训练配置文件;

data : 训练数据集配置文件;

epochs : 训练迭代轮数;

batch-size :每次迭代送入神经网络进行训练的图片数量,需为16倍数,根据GPU显存设置,越大越耗显存,训练速度越快;

imgsz : 训练指定输入图片的尺寸,所有输入图片在送入模型前都会被resize成指定的大小,尺寸大检测小目标效果好,训练慢,反之精度差,训练快;

patience :早停机制,迭代多少轮验证指标(如val_lossmAP)不提升,就提前停止训练,以防止过拟合、节省资源;

2.7 开始训练

执行python train.py

可能报以下错误:

  1. All git commands will error until this is rectified错误提示

按照提示执行export GIT_PYTHON_REFRESH=quiet继续执行训练命令,就可以开始训练了。

  1. 卡在not pretrained!!!!!!!!!!!!!!!!!!!!

如果误将data/images/自带的两张测试图片删除,会卡在not pretrained!!!!!!!!!!!!!!!!!!!!后,AMP: checks passed前,解决方式可以将图片下载下来重新放回,也可以修改utils/general.py按照如下修改:

im = f if f.exists() else 'https://ultralytics.com/images/bus.jpg' if False else np.ones((640, 640, 3))
# im = f if f.exists() else 'https://ultralytics.com/images/bus.jpg' if check_online() else np.ones((640, 640, 3))
  1. 下载Arial.ttf 字体错误

手动下载Arial.ttf 字体文件,放到docker环境下的/root/.config/Ultralytics/文件夹下;

  1. 获取git信息错误

删除.git文件夹

2.8 验证训练结果

训练结束后在代码根目录下执行检测命令,可以将待检测图片放到data/samples目录下执行

python detect.py --weights runs/train/exp/weights/best.pt --source data/samples/ --device 0 --data data/fall.yaml

注意: 每训练一次都会在runs/train/目录下新创建一个exp加数字文件夹,运行测试用例时选择最新的,测试结果也会保存在runs/detect目录下最新的exp文件夹下

2.9 导出模型

以导出onnx为例,执行以下命令:

# 安装onnx支持
pip3 install onnxsim -i http://mirrors.aliyun.com/pypi/simple/ --trusted-host mirrors.aliyun.com
# 导出[640,480]模型
# parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640, 640], help='image (h, w)') 尺寸 h在前
python export.py --weights runs/train/exp8/weights/best.pt --img-size 480 640 --batch 1 --device 0 --opset 12 --optimize --dynamic --include onnx --device cpu

3. 异常解决

3.1. Docker 守护进程无法找到具备 GPU 能力的设备驱动

docker: Error response from daemon: could not select device driver ““ with capabilities: [[gpu]].
  1. 检查主机 GPU 和 NVIDIA 驱动是否正常工作

nvidia-smi
  1. 检查 NVIDIA 容器工具包是否安装

dpkg -l | grep nvidia-container-toolkit

如果没有任何信息,则使用以下命令安装:

distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
    && curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - \
    && curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt update
sudo apt install -y nvidia-container-toolkit
sudo systemctl restart docker
  1. 为Docker添加NVIDIA配置

修改/etc/docker/daemon.json 文件追加以下内容:

{
	"default-runtime": "nvidia",
	"runtimes": {
		"nvidia": {
			"path": "nvidia-container-runtime",
			"runtimeArgs": []
		}
	}
}

保存后重启sudo systemctl restart docker

3.1. RuntimeError: Numpy is not available

这是因为Numpy版本太高,将现有Numpy卸载

pip uninstall numpy

安装numpy=1.26.4,解决此问题

pip install numpy==1.26.4 -i https://pypi.tuna.tsinghua.edu.cn/simple

0

评论区