YOLO_V5(Pytorch版本)模型训练以及rknn模型(rknn-toolkit平台)转换部署和测试流程(Linux)

• 说明：本文章用于学习与指导开发者，如何使用Pytorch版本的YOLO V5框架训练自己的模型、如何使用rknn-toolkit模型转换平台进行模型转换以及使用rk1808实现模型转换以及模型的测试

• 关键词：YOLO V5; RKNN; 模型转换; NPU推理计算

1.使用YOLO V5训练自己的数据集

1.1 安装Miniconda

• 可参照本人博客：点击链接直达 步骤一和步骤二

1.2 安装Tensorflow、Pytorch(不用提前安装)

• 在这里推荐使用pip安装
• 推荐使用源安装
• • 豆瓣源：https://pypi.tuna.tsinghua.edu.cn/simple
• • 清华源：https://pypi.douban.com/simple/

安装代码：

pip install +包名(==版本号) -i +源名称

1.3 下载YOLO V5官方版本

• 下载地址：点击此处直达

• 到达该链接后，点击克隆->下载.zip文件，千万不要使用git clone命令，否则下的不是这个版本（亲测非该版本模型转换时候会出错）

1.4 部署相关文件完成训练自己的数据集

• 下载后，进入pytorch环境进入yolov5文件夹，使用换源的方法安装依赖（注意参考1.2）。

注：如果你前面安装时没有换源，我强烈建议你使用换源的方法在安装一次;安装过的模块不会在安装，以防缺少模块，影响后续程序运行以及模型训练。

在新下载的YOLO V5文件夹下会有./requirements.txt文件，这里面包含了训练时所需要安装的相关依赖，我们需要将其中的>=改为==这样能够保证各个依赖的版本是相互匹配的。

安装命令：

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

清华源安装速度如果较慢，则可用豆瓣源,具体命令如下：

pip install -r requirements.txt -i https://pypi.douban.com/simple/

1.4.1 准备工作

• 在 yolov5目录下[新建文件夹]VOCData（可以自定义命名）
  

• 在VOCData下[新建两个文件夹]Annotations以images
  

注：Annotations文件：用于存放标注图片后产生的内容（这里采用XML格式）,即标签的信息；images文件：用于存放要标注的图片（jpg格式）,即图片的信息

• 模型标注工具LabelImg安装
  可参考：点击链接直达

• LabelImg使用
  可参考：点击链接直达

1.4.2 划分训练集、验证集、测试集

• 在VOCData目录下创建程序 split_train_val.py 并运行

# coding:utf-8

import os
import random
import argparse

parser = argparse.ArgumentParser()
#xml文件的地址，根据自己的数据进行修改 xml一般存放在Annotations下
parser.add_argument('--xml_path', default='VOCData/Annotations', type=str, help='input xml label path')
#数据集的划分，地址选择自己数据下的ImageSets/Main
parser.add_argument('--txt_path', default='VOCData/ImageSets/Main', type=str, help='output txt label path')
opt = parser.parse_args()

trainval_percent = 1.0  # 训练集和验证集所占比例。 这里没有划分测试集
train_percent = 0.9     # 训练集所占比例，可自己进行调整
xmlfilepath = opt.xml_path
txtsavepath = opt.txt_path
total_xml = os.listdir(xmlfilepath)
if not os.path.exists(txtsavepath):
    os.makedirs(txtsavepath)

num = len(total_xml)
list_index = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list_index, tv)
train = random.sample(trainval, tr)

file_trainval = open(txtsavepath + '/trainval.txt', 'w')
file_test = open(txtsavepath + '/test.txt', 'w')
file_train = open(txtsavepath + '/train.txt', 'w')
file_val = open(txtsavepath + '/val.txt', 'w')

for i in list_index:
    name = total_xml[i][:-4] + '
'
    if i in trainval:
        file_trainval.write(name)
        if i in train:
            file_train.write(name)
        else:
            file_val.write(name)
    else:
        file_test.write(name)

file_trainval.close()
file_train.close()
file_val.close()
file_test.close()

注意：修改xml_path和txt_path中的默认参数（不修改也没有关系），以及trainval_percent用于划分测试集，为1.0则表示没有划分测试集

运行完毕后 会生成 ImagesSets/Main 文件夹，且在其下生成 测试集、训练集、验证集，存放图片的名字（无后缀.jpg）

1.4.3 XML格式转yolo_txt格式

在VOCData目录下创建程序 xml_to_yolo.py 并运行

# -*- coding: utf-8 -*-
import xml.etree.ElementTree as ET
import os
from os import getcwd

sets = ['train', 'val', 'test']
classes = ["head", "body", "lifejacket"]   # 改成自己的类别
abs_path = os.getcwd()
print(abs_path)

def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    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):
    in_file = open('/home/lch/yolo/yolov5-master/VOCData/Annotations/%s.xml' % (image_id), encoding='UTF-8')
    out_file = open('/home/lch/yolo/yolov5-master/VOCData/labels/%s.txt' % (image_id), 'w')
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)
    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        #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))
        b1, b2, b3, b4 = b
        # 标注越界修正
        if b2 > w:
            b2 = w
        if b4 > h:
            b4 = h
        b = (b1, b2, b3, b4)
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '
')

wd = getcwd()
for image_set in sets:
    if not os.path.exists('/home/lch/yolo/yolov5-master/VOCData/labels/'):
        os.makedirs('/home/lch/yolo/yolov5-master/VOCData/labels/')
    image_ids = open('/home/lch/yolo/yolov5-master/VOCData/ImageSets/Main/%s.txt' % (image_set)).read().strip().split()
   
    if not os.path.exists('/home/lch/yolo/yolov5-master/VOCData/dataSet_path/'):
        os.makedirs('/home/lch/yolo/yolov5-master/VOCData/dataSet_path/')
     
    list_file = open('VOCData/dataSet_path/%s.txt' % (image_set), 'w')
    # 这行路径不需更改，这是相对路径
    for image_id in image_ids:
        list_file.write('/home/lch/yolo/yolov5-master/VOCData/images/%s.jpg
' % (image_id))
        convert_annotation(image_id)
    list_file.close()

注：需要将第 7 行改成自己所标注的类别以及修改各绝对路径

运行后会生成如下 labels 文件夹和 dataSet_path 文件夹。其中labels中为不同图像的标注文件。每个图像对应一个txt文件，文件每一行为一个目标的信息，包括class, x_center, y_center, width, height格式，这种即为 yolo_txt格式；dataSet_path文件夹包含三个数据集的txt文件，train.txt等txt文件为划分后图像所在位置的路径，如train.txt就含有所有训练集图像的路径。

1.4.4 修改配置文件

在 yolov5 目录下的 data 文件夹下 新建一个 myvoc.yaml文件（可以自定义命名）。内容为：训练集以及验证集（train.txt和val.txt）的路径（可以改为相对路径）以及 目标的类别数目和类别名称。

1.4.5 修改参数

修改models/yolov5s.yaml,只需更改nc(标注类别数)，不用更改anchors

1.5 开始训练

1.5.1 修改train.py中的默认参数

• weights：不用修改，起初使用yolov5s.pt
• cfg:需改为yolov5s.yaml所在地址 -> models/yolov5s.yaml
• data:需要修改为myvoc.yaml所在地址 -> data/myvoc.yaml
• epochs：指的就是训练过程中整个数据集将被迭代（训练）了多少次，显卡不行你就调小点(根据需要修改)。
• batch-size：训练完多少张图片才进行权重更新，显卡不行就调小点(本人修改为8)。
• img-size：输入图片宽高，显卡不行就调小点（YOLO V5一般设置为640，支持默认即可）。
• device：cuda device, i.e. 0 or 0,1,2,3 or cpu。选择使用GPU还是CPU[device:0->GPU(速度快)；device：cpu->CPU(速度慢)]使用GPU需要显卡支持以及安装显卡驱动。

1.5.2 训练运行方式 (案例，需要具体问题具体分析）

python train.py

或者

python train.py --weights weights/yolov5s.pt  --cfg models/yolov5s.yaml  --data data/myvoc.yaml --epoch 100 --batch-size 8 --img 640   --device 0

1.6 训练结束后得到模型

• 模型存放地址：runs/train/exp/weights。在该文件夹下会有两个.pt文件，这两个文件分别是best.pt和last.pt文件，分别存放的是训练过程中最佳和最后的模型文件。
  

• 新建文件夹，用于存放这两个模型文件

1.7 测试训练后的pytorch版本模型(.pt文件)

• 打开主目录文件夹下的detect.py文件

1.7.1 修改detect.py中的默认参数

• weights：需要修改，默认的为yolov5s.pt，需要将该参数指向训练完成后的模型文件****/last.pt
• source:待检测图片的存放位置(可以存放多张图片)
• imgsz:图片大小设置，需要与train.py中的img-size参数设置相同（因为YOLO V5图片输入有相应的要求，因此默认设置为640，不需要更改）
• device:与train.py中的device参数保持一致。

1.7.2 检测运行方式（案例，需要具体问题具体分析）

python detect.py

或者

python detect.py --weights runs/train/exp4/weights/best.pt --source VOCData/images/0.jpg

注：如果检测没有发生报错，且能够在runs/detect文件中找到相应的检测结果，说明文件训练没有问题，模型能够正常使用

1.7.3 检测案例展示

• 原始图片
  

• 检测后的图片
  

• 到这里为止，YOLO V5训练自己的模型得到pytorch模型告一段落

1.8 运行export.py文件

1.8.1 修改export.py代码

• 1.需要修改 yolov5/models/yolo.py文件的后处理部分，将class Detect(nn.Module) 类的子函数forward由

  def forward(self, x):
          z = []  # inference output
          for i in range(self.nl):
              x[i] = self.m[i](x[i])  # conv
              bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
              x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
  
              if not self.training:  # inference
                  if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:
                      self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)
  
                  y = x[i].sigmoid()
                  if self.inplace:
                      y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
                      y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
                  else:  # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
                      xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
                      wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
                      y = torch.cat((xy, wh, y[..., 4:]), -1)
                  z.append(y.view(bs, -1, self.no))
  
          return x if self.training else (torch.cat(z, 1), x)
  

  修改为：

  def forward(self, x):
          z = []  # inference output
          for i in range(self.nl):
              x[i] = self.m[i](x[i])  # conv
  
          return x
  

• 2.修改export.py中的默认参数

• • weights：需要修改，默认的为yolov5s.pt，需要将该参数指向训练完成后的模型文件****/last.pt
• • mg-size：YOLO V5默认为640
• • device：因为使用GPU加速，因此设置为0

1.8.2 运行export.py

python export.py

或者

python3 export.py --weight finial_model/last.pt  --img 640 --batch 1 --include torchscript

得到新的.pt文件 -> finial_model/last.torchscript.pt

注意：遇到No model ***问题，可以采用pip install ***方法解决

---

2.rknn-toolkit模型转换工具安装

可参照我的博客：点击链接直达

---

3.模型转换

3.1 下载并解压rknn-toolkit-master

3.2 进入目录

cd ~/rknn-toolkit-master/examples/pytorch/yolov5

3.3 复制模型文件

• 将新的.pt文件 -> finial_model/last.torchscript.pt复制到目录文件夹下

• 替换测试图片
  

• 修改dataset.txt(将其中的路径修改为测试图片的路径)
  

3.4 修改test.py文件

3.5 注意事项

• 在模型转换过程中需要保证pytorch版本与训练时的朋友torch版本一致
• 需要安装torchvision，但要注意torchvision与pytorch的版本对应关系

3.6 运行模型转换程序

3.6.1 激活rknn-toolkit环境(自行创建完成，参照2.rknn-toolkit模型转换工具安装)

conda activate rknn

3.6.2 rknn模型生成并返回测试图片识别结果

python test.py

4.编写代码调用rknn模型完成识别

4.1 新建detect.py文件

import os
import numpy as np
import cv2
from rknn.api import RKNN

BOX_THRESH = 0.5
NMS_THRESH = 0.6
IMG_SIZE = (640, 640) # (width, height), such as (1280, 736)

CLASSES = ("head", "body", "lifejacket")

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def xywh2xyxy(x):
    # Convert [x, y, w, h] to [x1, y1, x2, y2]
    y = np.copy(x)
    y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
    y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
    y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
    y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
    return y

def process(input, mask, anchors):

    anchors = [anchors[i] for i in mask]
    grid_h, grid_w = map(int, input.shape[0:2])

    box_confidence = sigmoid(input[..., 4])
    box_confidence = np.expand_dims(box_confidence, axis=-1)

    box_class_probs = sigmoid(input[..., 5:])

    box_xy = sigmoid(input[..., :2])*2 - 0.5

    col = np.tile(np.arange(0, grid_w), grid_h).reshape(-1, grid_w)
    row = np.tile(np.arange(0, grid_h).reshape(-1, 1), grid_w)
    col = col.reshape(grid_h, grid_w, 1, 1).repeat(3, axis=-2)
    row = row.reshape(grid_h, grid_w, 1, 1).repeat(3, axis=-2)
    grid = np.concatenate((col, row), axis=-1)
    box_xy += grid
    box_xy *= (int(IMG_SIZE[1]/grid_h), int(IMG_SIZE[0]/grid_w))

    box_wh = pow(sigmoid(input[..., 2:4])*2, 2)
    box_wh = box_wh * anchors

    box = np.concatenate((box_xy, box_wh), axis=-1)

    return box, box_confidence, box_class_probs

def filter_boxes(boxes, box_confidences, box_class_probs):
    """Filter boxes with box threshold. It's a bit different with origin yolov5 post process!

    # Arguments
        boxes: ndarray, boxes of objects.
        box_confidences: ndarray, confidences of objects.
        box_class_probs: ndarray, class_probs of objects.

    # Returns
        boxes: ndarray, filtered boxes.
        classes: ndarray, classes for boxes.
        scores: ndarray, scores for boxes.
    """
    boxes = boxes.reshape(-1, 4)
    box_confidences = box_confidences.reshape(-1)
    box_class_probs = box_class_probs.reshape(-1, box_class_probs.shape[-1])

    _box_pos = np.where(box_confidences >= BOX_THRESH)
    boxes = boxes[_box_pos]
    box_confidences = box_confidences[_box_pos]
    box_class_probs = box_class_probs[_box_pos]

    class_max_score = np.max(box_class_probs, axis=-1)
    classes = np.argmax(box_class_probs, axis=-1)
    _class_pos = np.where(class_max_score* box_confidences >= BOX_THRESH)

    boxes = boxes[_class_pos]
    classes = classes[_class_pos]
    scores = (class_max_score* box_confidences)[_class_pos]

    return boxes, classes, scores

def nms_boxes(boxes, scores):
    """Suppress non-maximal boxes.

    # Arguments
        boxes: ndarray, boxes of objects.
        scores: ndarray, scores of objects.

    # Returns
        keep: ndarray, index of effective boxes.
    """
    x = boxes[:, 0]
    y = boxes[:, 1]
    w = boxes[:, 2] - boxes[:, 0]
    h = boxes[:, 3] - boxes[:, 1]

    areas = w * h
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)

        xx1 = np.maximum(x[i], x[order[1:]])
        yy1 = np.maximum(y[i], y[order[1:]])
        xx2 = np.minimum(x[i] + w[i], x[order[1:]] + w[order[1:]])
        yy2 = np.minimum(y[i] + h[i], y[order[1:]] + h[order[1:]])

        w1 = np.maximum(0.0, xx2 - xx1 + 0.00001)
        h1 = np.maximum(0.0, yy2 - yy1 + 0.00001)
        inter = w1 * h1

        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        inds = np.where(ovr <= NMS_THRESH)[0]
        order = order[inds + 1]
    keep = np.array(keep)
    return keep


def yolov5_post_process(input_data):
    masks = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
    anchors = [[10, 13], [16, 30], [33, 23], [30, 61], [62, 45],
              [59, 119], [116, 90], [156, 198], [373, 326]]

    boxes, classes, scores = [], [], []
    for input,mask in zip(input_data, masks):
        b, c, s = process(input, mask, anchors)
        b, c, s = filter_boxes(b, c, s)
        boxes.append(b)
        classes.append(c)
        scores.append(s)

    boxes = np.concatenate(boxes)
    boxes = xywh2xyxy(boxes)
    classes = np.concatenate(classes)
    scores = np.concatenate(scores)

    nboxes, nclasses, nscores = [], [], []
    for c in set(classes):
        inds = np.where(classes == c)
        b = boxes[inds]
        c = classes[inds]
        s = scores[inds]

        keep = nms_boxes(b, s)

        nboxes.append(b[keep])
        nclasses.append(c[keep])
        nscores.append(s[keep])

    if not nclasses and not nscores:
        return None, None, None

    boxes = np.concatenate(nboxes)
    classes = np.concatenate(nclasses)
    scores = np.concatenate(nscores)

    return boxes, classes, scores

def draw(image, boxes, scores, classes):
    """Draw the boxes on the image.
    # Argument:
        image: original image.
        boxes: ndarray, boxes of objects.
        classes: ndarray, classes of objects.
        scores: ndarray, scores of objects.
        all_classes: all classes name.
    """
    for box, score, cl in zip(boxes, scores, classes):
        left, top, right, bottom = box
        print('class: {}, score: {}'.format(CLASSES[cl], score))
        print('box coordinate left,top,right,bottom: [{}, {}, {}, {}]'.format(left, top, right, bottom))
        left = int(left)
        top = int(top)
        right = int(right)
        bottom = int(bottom)

        cv2.rectangle(image, (left, top), (right, bottom), (255, 0, 0), 2)
        cv2.putText(image, '{0} {1:.2f}'.format(CLASSES[cl], score),
                    (left, top - 6),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.6, (0, 0, 255), 2)

if __name__=="__main__":

    # path
    rknn_model_path="last.rknn"
    img_path="test_picture.jpg"

    # Create RKNN object
    rknn = RKNN(verbose=True)

    rknn.load_rknn(rknn_model_path)

    ret = rknn.init_runtime(target='rk1808',device_id='TS0180842112009A5')

    # 读取图片
    # 注意需要将图片转成RGB格式
    frame = cv2.imread(img_path)
    image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    image = cv2.resize(image, (640,640))

    # 进行推断
    # result = rknn.inference(inputs=[image])

    print('--> Running model')
    outputs = rknn.inference(inputs=[image])

    # post process
    input0_data = outputs[0]
    input1_data = outputs[1]
    input2_data = outputs[2]

    input0_data = input0_data.reshape([3,-1]+list(input0_data.shape[-2:]))
    input1_data = input1_data.reshape([3,-1]+list(input1_data.shape[-2:]))
    input2_data = input2_data.reshape([3,-1]+list(input2_data.shape[-2:]))

    input_data = list()
    input_data.append(np.transpose(input0_data, (2, 3, 0, 1)))
    input_data.append(np.transpose(input1_data, (2, 3, 0, 1)))
    input_data.append(np.transpose(input2_data, (2, 3, 0, 1)))

    boxes, classes, scores = yolov5_post_process(input_data)

    img_1 = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    if boxes is not None:
        draw(img_1, boxes, scores, classes)
    cv2.imshow("post process result", img_1)
    cv2.waitKeyEx(0)

    rknn.release()

4.2 运行detect.py文件

python detect.py

完成！！！！！！！！！！！！！！！！