零基础：训练并部署 TensorFlow 版 YOLOv5 模型

概述

目标检测是计算机视觉上的一个重要任务，本文介绍的是 YOLO算法，其全称是 You Only Look Once: Unified, Real-Time Object Detection，它是目标检测中实现端到端目标检测的佼佼者，从YOLOv1 到 YOLOv5 其中在数据处理、网络结构上都做了不少优化，而 YOLOv5 能够达到体积更小、精度更好，本文就从零开始介绍如何通过用 TensorFlow 对 YOLOv5 进行搭建训练和部署。本实例源码可在点击以下链接：https://github.com/Yunying-CN/Yolov5-TF。

1.安装 Tensorflow2.x

为了提高训练速度减少训练时长，在训练阶段最好在配有 GPU的本地服务器或者云服务器上进行。本例以 Linux 64 位下的Python 3.8 版本为例，可选择下载对应的安装包。在保存安装包的路径下打开终端，运行命令进行安装 TensorFlow。这里安装的是 Tensorflow2.3.0-gpu 版本，搭配 cuda10.1 和对应的 cudnn，也可以直接通过 pip 安装命令来下载安装，如果速度较慢可以修改下载的源。

$sudo apt-get install python-pip python-dev
$pip3 install --upgrade pip
$pip3 install tensorflow-gpu==2.3.1 -i https://pypi.tuna.tsinghua.edu.cn/simple

安装完成后，可以终端打开 python 并导入 Tensorflow 来查看版本来验证是否安装成功。

$python3>>import tensorflow as tf>>tf.version

2.训练 YOLOv5 网络

YOLO 是经典的目标检测识别网络，而 YOLOv5 是 YOLO 系列中识别率最高而且速度最快的目标检测识别模型。YOLOv5模型属于监督学习，训练模型的样本需要包括物体的位置坐标（矩形框）和物体所属的类别。将数据集中的图片作为网络输入，物体的类别和坐标作为标签信息，对网络进行训练，得到网络对物体检测和识别的能力。

2.1 数据集

本实例以开源的 Pascal Voc2012 数据集。Pascal VOC2012作为基准数据之一，在对象检测、图像分割网络对比实验与模型效果评估中被频频使用，Pascal VOC2012 数据集主要是针对视觉任务中监督学习提供标签数据，它一共包含有 20 个类别，分别为：aeroplane、bicycle、bird、boat、bottle、bus、car、cat、chair、cow、dining table、dog、horse、motorbike、person、potted plant、sheep、sofa、train、tv/monitor，训练图像有 5717 张，目标数 13609 个，测试 图 像 有 11540 张， 目 标 数 27450 个。Pascal Voc2012数据集可以在官网上进行下载 (http://host.robots.ox.ac.uk/pascal/VOC/voc2012/)，也可以在终端通过命令下载数据集并解压。

$wget http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar -O ./data/voc2012.tar
$mkdir -p ./data/voc
$tar -xf ./data/voc2012.tar -C ./data/voc
$ls ./data/voc

里面包括有 Annotations、ImageSets、JPEGImages、Segmenta-tionClass SegmentationObject 五个文件夹，Annotations 文件夹中保存了进行目标检测任务时的标签文件为 .xml 格式，标签文件名与图片名一一对应。.xml 文档记录了该图片的尺寸信息以及图片中识别物体的类别和其具体位置信息。ImageSets包含三个子文件夹 Layout、Main、Segmentation，其中 Main 存放的是分类和检测的数据集分割文件，JPEGImages 存放 .jpg 格式的图片文件，SegmentationClass 存放按照类别进行分割的图片，SegmentationObject 存放按照物体进行分割的图片。

如果想制作自己的数据集，也可以通过使用 LabelImg 工具，框选出图像中所需识别物体的位置和标记该物体的类别并将所有所识别的类型保存在 .names 文件中。

2.2 生成 TFRecord 文件

为了高效的读取数据，可以将数据进行序列化存储，便于网络流式读取数据，TFRecord 是存储二进制序列数据的文件格式，只占用一个内存块，保存记录的方法可以任意的数据转换为TensorFlow 所支持的格式，这种方法可以使 TensorFlow 的数据集更容易与网络应用架构相匹配。在生成 TFRecord 文件中需要保存训练图片及图片中物体的类别和位置信息。

import os
from absl import app, flags, logging
from absl.flags import FLAGS
import tensorflow as tf
import lxml.etree
import tqdm

flags.DEFINE_string(‘data_dir’, ‘./data/voc/VOCdevkit/VOC2012/’,
’path to PASCAL VOC dataset’)
flags.DEFINE_enum(‘split’, ‘val’, [
‘train’, ‘val’], ‘specify train or val spit’)
flags.DEFINE_string(‘output_file’, ‘./data/voc2012_train.tfrecord’, ‘outpot dataset’)
flags.DEFINE_string(‘classes’, ‘./data/voc2012.names’, ‘classes file’)

def build_example(annotation, class_map):img_path = os.path.join(
FLAGS.data_dir, ‘JPEGImages’, annotation[‘filename’])
img_raw = open(img_path, ‘rb’).read()
width = int(annotation[‘size’][‘width’])
height = int(annotation['size’][‘height’])
xmin = []
ymin = []
xmax = []
ymax = []
classes = []
classes_text = []
if ‘object’ in annotation:
for obj in annotation[‘object’]:
xmin.append(float(obj[‘bndbox’][‘xmin’]) / width)
ymin.append(float(obj[‘bndbox’]['ymin’]) / height)
xmax.append(float(obj[‘bndbox’][‘xmax’]) / width)
ymax.append(float(obj[‘bndbox’][‘ymax’]) / height)
classes_text.append(obj[‘name’].encode(‘utf8’))
classes.append(class_map[obj[‘name’]])
example = tf.train.Example(features=tf.train.Features(feature={‘image/encoded’: tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw])),
’image/object/bbox/xmin’: tf.train.Feature(float_list=tf.train.FloatList(value=xmin)),
’image/object/bbox/xmax’: tf.train.Feature(float_list=tf.train.FloatList(value=xmax)),
’image/object/bbox/ymin’: tf.train.Feature(float_list=tf.train.FloatList(value=ymin)),
’image/object/bbox/ymax’: tf.train.Feature(float_list=tf.train.FloatList(value=ymax)),
’image/object/class/text’: tf.train.Feature(bytes_list=tf.train.BytesList(value=classes_text)),
’image/object/class/label’: tf.train.Feature(int64_list=tf.train.Int64List(value=classes)),}))return example
def parse_xml(xml):
if not len(xml):
return {xml.tag: xml.text}
result = {}
for child in xml:
child_result = parse_xml(child)
if child.tag != ‘object’:
result[child.tag] = child_result[child.tag]
else:
if child.tag not in result:
result[child.tag] =[]
result[child.tag].append(child_result[child.tag])
return {xml.tag: result}
def main(_argv):
class_map = {name: idx for idx, name in enumerate(open(FLAGS.classes).read().splitlines())}
writer = tf.io.TFRecordWriter(FLAGS.output_file)
image_list = open(os.path.join(
FLAGS.data_dir, ‘ImageSets’, ‘Main’, ‘%s.txt’ % FLAGS.split)).
read().splitlines()
logging.info(“Image list loaded: %d”, len(image_list))for name in tqdm.tqdm(image_list):
annotation_xml = os.path.join(FLAGS.data_dir, ‘Annotations’, name + ‘.xml’)
annotation_xml = lxml.etree.fromstring(open(annotation_xml).read())
annotation = parse_xml(annotation_xml)[‘annotation’]tf_example = build_example(annotation, class_map)writer.write(tf_example.SerializeToString())
writer.close()
logging.info(“Done”)

if name == ‘main’:
app.run(main)

2.3 YOLOv5 网络结构

YOLOv5 目标检测网络中一共有 4 个版本，分别是 YOLOv5s、YOLOv5m、YOLOv5l、YOLOv5x 四个模型，通过用 .yaml文件来配置模型。在 yaml 文件里面分别定义了各个参数变量如 : nc 代表分类目标的数量，depth_multiple 即网络深度表示 channel 的缩放系数，即将配置里面的 backbone和 head 部分有关通道的设置。而 width_multiple 即网络宽度表示BottleneckCSP模块的层缩放系数，将所有的BottleneckCSP 模块的 number 系数乘上该参数即为最终的层个数。通过这参数就可以实现不同大小不同复杂度的模型设计，4 个版本的 YOLOv5 也做了不同的设计。Anchors 为预设锚定框，预设了 640 × 640 图像大小下 9 种锚定框的尺寸。此外还有模型的主干网络backbone和通用检测层head，head 主要用于最终检测部分。它在特征图上应用锚定框并生成带有类概率、对象得分和边界框的最终输出向量。以下是以YOLOv5s.yaml 为例。

#Parameters
nc: 20 # number of classes
depth_multiple: 0.67 # model depth multiple
width_multiple: 0.75 # layer channel multiple
anchors:
-[10,13, 16,30, 33,23] # P3/8-
-[30,61, 62,45, 59,119] # P4/16
-[116,90, 156,198, 373,326] # P5/32
#YOLOv5 backbone
backbone:
#[from, number, module, args]
[[-1, 1, Focus, [64, 3]], # 0-P1/2
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
[-1, 3, C3, [128]],
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
[-1, 9, C3, [256]],
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
[-1, 9, C3, [512]],
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
[-1, 1, SPP, [1024, [5, 9, 13]]],
[-1, 3, C3, [1024, False]], # 9
]
#YOLOv5 head
head:
[[-1, 1, Conv, [512, 1, 1]],
[-1, 1, Upsample, [None, 2, ‘nearest’]],
[[-1, 6], 1, Concat, [-1]], # cat backbone P4
[-1, 3, C3, [512, False]], # 13
[-1, 1, Conv, [256, 1, 1]],
[-1, 1, Upsample, [None, 2, ‘nearest’]],
[[-1, 4], 1, Concat, [-1]], # cat backbone P3
[-1, 3, C3, [256, False]], # 17 (P3/8-small)
[-1, 1, Conv, [256, 3, 2]],
[[-1, 14], 1, Concat, [-1]], # cat head P4
[-1, 3, C3, [512, False]], # 20 (P4/16-medium)
[-1, 1, Conv, [512, 3, 2]],
[[-1, 10], 1, Concat, [-1]], # cat head P5
[-1, 3, C3, [1024, False]], # 23 (P5/32-large)
[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
]

2.4 YOLOv5 网络模块

YOLOv5 网络可以分为输入图像、backbone 主干网络以获得图像的特征图、head 检测头用作预测目标物体和位置，neck为对特征图在输入 head 前的特殊处理，4 个部分组成整个网络。而这 4 部分由几个不同的模块堆叠得到。

2.4.1 CBL 模块

CBL 为卷积模块，YOLOv5 主干网络中的 CBL 模块以 Convolu-tion + Batch Normalization + Activation 的形式，对输入数据进行卷积计算、批标准化计算和经过一个激活函数，其中的激活函数选用 LeakyRelu，对网络加入非线性并加快网络的收敛速度。

2.4.2 Focus 模块

从 .yaml 配置文件中可以看到在 backbone 主干网络中包含了focus 模块，focus 模块是对图片进行切片操作，通过在图片中每间隔 1 个像素取值，得到 4 张图片，使得图片的长和宽分别减半，通道数扩展为原来的 4 倍，该操作类似于 2 倍下采样但是保证了图片信息没有丢失，以 YOLOV5s 为例，原始的 640 × 640 × 3 的图像通过 Focus 模块，输出得到 320 × 320 × 12 的特征图。

2.4.3 bottleneck 模块

Bottleneck模块可以通过卷积计算改变数据的通道数，bottleneck 瓶颈层有多种形式，其标准形式为进行一个 1 × 1和 3 × 3 的卷积后加上其本身的短路连接， 而 BottleneckCSP是几个标准 bottleneck 的堆叠，YOLOV5 网络中的 C3 模块与 BottleneckCSP 模块类似，只是在 C3 中的卷积计算后加上了 BN 层和激活函数积操作。

2.4.4 SPP 模块

在目标检测中，通常输入图像的尺寸大小并不固定，为了得到统一大小的特征图，从 YOLOv3 开始引入了 SPP 空间金字塔池化模块，通过使用 CBL 模块使其通道数减半，然后将输入经过三个不同尺寸大小的最大池化层，连同输入通过 concat级联在一起，最后通过 CBL 模块使通道数减半，保证不同大小的输入在池化后的特征图长和宽能保持一致。

2.4.5 Upsample 模块

在进行预测前，网络对都特征图做了两次向上采样，得到 3 个不同尺寸大小的特征图，使图像的长和宽分别扩展为原来的 2倍，这能够使得相对较大或较小的物体都能更好地识别。
import tensorflow as tf
from tensorflow.keras.layers import Layer, Conv2D, BatchNormalization, MaxPool2D
from tensorflow import keras
import math
import numpy as np

class Conv2d(keras.layers.Layer):
def init(self, c1, c2, k, s=1, g=1, bias=True, w=None):
super(Conv2d, self).init()
assert g == 1, "TF v2.2 Conv2D does not support ‘groups’ argument"
self.conv = keras.layers.Conv2D(
c2, k, s, ‘VALID’,use_bias=bias,
kernel_initializer=keras.initializers.Constant(w.weight.permute(2, 3, 1, 0).numpy()),
bias_initializer=keras.initializers.Constant(w.bias.numpy())
if bias else None )
def call(self, inputs):
return self.conv(inputs)

class LeakyRelu(object):
def call(self, x):
return tf.nn.leaky_relu(x)

class Conv(Layer):
def init(self, filters, kernel_size, strides, padding=‘SAME’, groups=1):
super(Conv, self).init()
self.conv = Conv2D(filters, kernel_size, strides, padding, groups=groups, use_bias=False,
kernel_initializer=tf.random_normal_initializer (stddev=0.01),
kernel_regularizer=tf.keras.regularizers.L2(5e-4))
self.bn = BatchNormalization()
self.activation = LeakyRelu()
def call(self, x):
return self.activation(self.bn(self.conv(x)))
self.conv = Conv2D(filters, kernel_size, strides, padding, groups=groups, use_bias=False,
kernel_initializer=tf.random_normal_initializer(stddev=0.01),kernel_regularizer=tf.keras.regularizers.L2(5e-4))
self.bn = BatchNormalization()
self.activation = LeakyRelu()
def call(self, x):
return self.activation(self.bn(self.conv(x)))

class Focus(Layer):
def init(self, filters, kernel_size, strides=1, padding=‘SAME’):
super(Focus, self).init()
self.conv = Conv(filters, kernel_size, strides, padding)def call(self, x):
return self.conv(tf.concat([x[…, ::2, ::2, :],
x[…, 1::2, ::2, :],
x[…, ::2, 1::2, :],
x[…, 1::2, 1::2, :]],
axis=-1))

class Bottleneck(Layer):
def init(self, units, shortcut=True, expansion=0.5):
super(Bottleneck, self).init()
self.conv1 = Conv(int(units * expansion), 1, 1)
self.conv2 = Conv(units, 3, 1)
self.shortcut = shortcut
def call(self, x):
if self.shortcut:
return x + self.conv2(self.conv1(x))
return self.conv2(self.conv1(x))

class BottleneckCSP(Layer):
def init(self, units, n_layer=1, shortcut=True, expansion=0.5):
super(BottleneckCSP, self).init()
units_e = int(units * expansion)
self.conv1 = Conv(units_e, 1, 1)
self.conv2 = Conv2D(units_e, 1, 1, use_bias=False, kernel_initializer=tf.random_normal_initializer(stddev=0.01))self.conv3 = Conv2D(units_e, 1, 1, use_bias=False, kernel_initializer=tf.random_normal_initializer(stddev=0.01))self.conv4 = Conv(units, 1, 1)
self.bn =BatchNormalization(momentum=0.03)
self.activation = LeakyRelu()
self.modules = tf.keras.Sequential([Bottleneck(units_e, shortcut, expansion=1.0) for _ in range(n_layer)])
class BottleneckCSP(Layer):
def init(self, units, n_layer=1, shortcut=True, expansion=0.5):
super(BottleneckCSP, self).init()
units_e = int(units * expansion)
self.conv1 = Conv(units_e, 1, 1)
self.conv2 = Conv2D(units_e, 1, 1, use_bias=False, kernel_initializer=tf.random_normal_initializer(stddev=0.01))self.conv3 = Conv2D(units_e, 1, 1, use_bias=False, kernel_initializer=tf.random_normal_initializer(stddev=0.01))self.conv4 = Conv(units, 1, 1)
self.bn = BatchNormalization(momentum=0.03)
self.activation = LeakyRelu()
self.modules = tf.keras.Sequential([Bottleneck(units_e, shortcut, expansion=1.0) for _ in range(n_layer)])
def call(self, x):
y1 = self.conv3(self.modules(self.conv1(x)))
y2 = self.conv2(x)
return self.conv4(self.activation(self.bn(tf.concat([y1, y2], axis=-1))))

class SPP(Layer):
def init(self, units, kernels=(5, 9, 13)):
super(SPP, self).init()
units_e = units // 2 # Todo:
self.conv1 = Conv(units_e, 1, 1)
self.conv2 = Conv(units, 1, 1)
self.modules = [MaxPool2D(pool_size=x, strides=1, padding=‘SAME’) for x in kernels]
def call(self, x):
x = self.conv1(x)
return self.conv2(tf.concat([x] + [module(x) for module in self.modules], axis=-1))

class SPPCSP(Layer):
#Cross Stage Partial Networks
def init(self, units, n=1, shortcut=False, expansion=0.5, kernels=(5, 9, 13)):
super(SPPCSP, self).init()
units_e = int(2 * units * expansion)
self.conv1 = Conv(units_e, 1, 1)
self.conv2 = Conv2D(units_e, 1, 1, use_bias=False, kernel_initializer=tf.random_normal_initializer(stddev=0.01))self.conv3 = Conv(units_e, 3, 1)
self.conv4 = Conv(units_e, 1, 1)
self.modules = [MaxPool2D(pool_size=x, strides=1, padding=‘same’) for x in kernels]
self.conv5 = Conv(units_e, 1, 1)
self.conv6 = Conv(units_e, 3, 1)
self.bn = BatchNormalization()self.act = LeakyRelu()self.conv7 = Conv(units, 1, 1)
def call(self, x):
x1 = self.conv4(self.conv3(self.conv1(x)))
y1 = self.conv6(self.conv5(tf.concat([x1] + [module(x1) for module in self.modules], axis=-1)))
y2 = self.conv2(x)
return self.conv7(self.act(self.bn(tf.concat([y1, y2], axis=-1))))
class Upsample(Layer):
def init(self, i=None, ratio=2, method=‘bilinear’):
super(Upsample, self).init()
self.ratio = ratio
self.method = method

def call(self, x):
return tf.image.resize(x, (tf.shape(x)[1] * self.ratio, tf.shape(x)[2] * self.ratio), method=self.method)

2.5 网络和损失函数

读取 .yaml 文件中的 backbone 和 head 结构可以以序列形式把上述定义好的网络模块堆叠起来完成网络框架。

在损失计算中，分类任务和置信度任务都是通过二元交叉熵损失函数计算，再通过 gamma 和 alpha 的 Focal Loss 来调整权重，而边界框是通过以 GIOU 来计算其损失函数。

def parse_model(yaml_dict): # model_dict, input_
channels(3)
anchors, nc = yaml_dict[‘anchors’], yaml_dict[‘nc’]
depth_multiple, width_multiple = yaml_dict[‘depth_multiple’], yaml_dict[‘width_multiple’]
num_anchors = (len(anchors[0]) // 2)
if isinstance(anchors, list) else anchorsoutput_dims = num_anchors * (nc + 5)
layers = []
#from, number, module, args
for i, (f, number, module, args) in enumerate(yaml_dict[‘backbone’] + yaml_dict[‘head’]):
#all component is a Class, initialize here, call in self.forward
module = eval(module) if isinstance(module, str) else module
for j, arg in enumerate(args):
try:
args[j] = eval(arg) if isinstance(arg, str) else arg except:
pass
number = max(round(number * depth_multiple), 1) if number > 1 else number
if module in [Conv2D, Conv, Bottleneck, SPP, Focus, BottleneckCSP, C3]:
c2 = args[0]
c2 = math.ceil(c2 * width_multiple / 8) * 8 if c2 != output_dims else c2
args = [c2, args[1:]]
if module in [BottleneckCSP, C3, SPPCSP]:
args.insert(1, number)
number = 1
modules = tf.keras.Sequential([module(*args) for _ in range(number)]) if number >
1 else module(*args)
modules.i, modules.f = i,
flayers.append(modules)
return layers

class Model(object):
#model, channels, classes
def init(self, cfg=‘yolov5s.yaml’, ch=3, nc=20, model=None, imgsz=(640, 640)):
super(Model, self).init()
if isinstance(cfg, dict):
self.yaml = cfg # model dict
else: # is *.yaml
import yaml # for torch hub
self.yaml_file = Path(cfg).name
with open(cfg) as f:
self.yaml = yaml.load(f, Loader=yaml.FullLoader) # model dict
self.imgsz =imgsz
#Define model
if nc and nc != self.yaml[‘nc’]:
print(‘Overriding %s nc=%g with nc=%g’% (cfg, self.yaml[‘nc’], nc))
self.yaml[‘nc’] = nc # override yaml valueself.model = parse_model(self.yaml)
if isinstance(model, Detect):
#transfer the anchors to grid coordinator, 3 * 3 * 2
model.anchors /= tf.reshape(module.stride, [-1, 1, 1])
def call(self, img_size, name=‘yolo’):
x = tf.keras.Input([img_size, img_size, 3])
output = self.forward(x)
return tf.keras.Model(inputs=x, outputs=output, name=name)
def forward(self, inputs, tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45, conf_thres=0.25):
y = [] # outputs
x = inputs
for i, m in enumerate(self.model):
if m.f != -1:
if isinstance(m.f, int):
x = y[m.f]
x = y[m.f]
else:
x = [x if j == -1 else y[j] for j in m.f]
x = m(x) # run
y.append(x)
return x

class Loss(object):
def init(self, anchors, iou_thres, num_classes=20, img_size=640, label_smoothing=0):
self.anchors = anchors
self.strides = [8, 16, 32]
self.iou_thres = iou_thres
self.num_classes = num_classes
self.img_size = img_size
self.bce_conf = tf.keras.losses.BinaryCrossentropy(reduction=tf.keras.losses.Reduction.NONE)
self.bce_class = tf.keras.losses.BinaryCrossentropy(reduction=tf.keras.losses.Reduction.NONE,label_smoothing=label_smoothing)
def call(self, y_true, y_pred):
iou_loss_all = obj_loss_all = class_loss_all = tf.zeros(1)
balance = [4.0, 1.0, 0.4] if len(y_pred) == 3 else [4.0, 1.0, 0.25, 0.06]
for i, (pred, true) in enumerate(zip(y_pred, y_true)):
true_box, true_obj, true_class = tf.split(true, (4, 1, -1),

axis=-1)
pred_box, pred_obj, pred_class = tf.split(pred, (4, 1, -1), axis=-1)
if tf.shape(true_class)[-1] == 1 and self.num_classes > 1:
true_class = tf.squeeze(tf.one_hot(tf.cast(true_class, tf.dtypes.int32),
depth=self.num_classes, axis=-1), -2)
box_scale = 2 - 1.0 * true_box[…, 2] * true_box[…, 3] / (self.img_size ** 2)
obj_mask = tf.squeeze(true_obj, -1) # obj or noobj
background_mask = 1.0 - obj_mask
conf_focal = tf.squeeze(tf.math.pow(true_obj - pred_obj, 2), -1)
#giou loss
iou = bbox_iou(pred_box, true_box, xyxy=False, giou=True) iou_loss = (1 - iou) * obj_mask * box_scale # batch_size * grid * grid * 3
#confidence loss
conf_loss = self.bce_conf(true_obj, pred_obj)
conf_loss = conf_focal * (obj_mask * conf_loss + background_mask * conf_loss)
#class loss
class_loss = obj_mask * self.bce_class(true_class, pred_class)
iou_loss = tf.reduce_mean(tf.reduce_sum(iou_loss, axis=[1, 2, 3]))
conf_loss = tf.reduce_mean(tf.reduce_sum(conf_loss, axis=[1, 2, 3]))class_loss = tf.reduce_mean(tf.reduce_sum(class_loss, axis=[1, 2, 3]))
iou_loss_all += iou_loss * balance[i]
obj_loss_all += conf_loss * balance[i]
class_loss_all += class_loss * self.num_classes * balance[i] # to balance the 3 loss

return (iou_loss_all, obj_loss_all, class_loss_all)

def bbox_iou(bbox1, bbox2, xyxy=False, giou=False, diou=False, ciou=False, epsilon=1e-9):
assert bbox1.shape == bbox2.shape
#giou loss: https://arxiv.org/abs/1902.09630
if xyxy:
b1x1, b1y1, b1x2, b1y2 = bbox1[…, 0], bbox1[…, 1], bbox1[…, 2], bbox1[…, 3]
b2x1, b2y1, b2x2, b2y2 = bbox2[…, 0], bbox2[…, 1], bbox2[…, 2], bbox2[…, 3]
else: # xywh -> xyxy
b1x1, b1x2 = bbox1[…, 0] - bbox1[…, 2] / 2, bbox1[…, 0] + bbox1[…, 2] / 2
b1y1, b1y2 = bbox1[…, 1] - bbox1[…, 3] / 2, bbox1[…, 1] + bbox1[…, 3] / 2
b2x1, b2x2 = bbox2[…, 0] - bbox2[…, 2] / 2, bbox2[…, 0] + bbox2[…, 2] / 2
b2y1, b2y2 = bbox2[…, 1] - bbox2[…, 3] / 2, bbox2[…, 1] + bbox2[…, 3] / 2

#intersection area
inter = tf.maximum(tf.minimum(b1x2, b2x2) - tf.maximum(b1x1, b2x1), 0) *
tf.maximum(tf.minimum(b1y2, b2y2) - tf.maximum(b1y1, b2y1), 0)

#union areaw1, h1 = b1x2 - b1x1 + epsilon, b1y2 - b1y1 + epsilonw2, h2 = b2x2 - b2x1+ epsilon, b2y2 - b2y1 + epsilonunion = w1 * h1 + w2 * h2 - inter + epsilon

#Giouiou = inter / union

cw = tf.maximum(b1x2, b2x2) - tf.minimum(b1x1, b2x1)ch = tf.maximum(b1y2, b2y2) - tf.minimum(b1y1, b2y1)enclose_area = cw * ch + epsilongiou = iou - 1.0 * (enclose_area - union) / enclose_areareturn tf.clip_by_value(giou, -1, 1)

2.6 传入训练数据设置训练参数

在完成网络的搭建后，需要从上述生成得到的 TFRecord 文件中读取训练数据，需要设置网络的分类类别数，根据 batch size 分批把数据放入网络中，并且设置网络训练轮数、优化器和学习率等，并将训练的网络模型保存为 .pb 或 .pbtxt 文件。

from absl import app, flags, logging
from absl.flags import FLAGS
import tensorflow as tf
import numpy as np
import cv2import time
from models.yolo import *
from data.dataset import *

flags.DEFINE_string(‘dataset’, ‘./data/voc2012_train.tfrecord’, ‘path to dataset’)
flags.DEFINE_string(‘val_dataset’, ‘./data/voc2012_val.tfrecord’, ‘path to validation dataset’)
flags.DEFINE_string(‘yaml_dir’, ‘./models/yolov5s.yaml’, ‘path to yaml file’)
flags.DEFINE_string(‘classes’, ‘./data/voc2012.names’, ‘path to classes file’)
flags.DEFINE_integer(‘epochs’, 20, ‘number of epochs’)
flags.DEFINE_integer(‘batch_size’, 8, ‘batch size’)
flags.DEFINE_integer(‘img_size’, 640, ‘image size’)
flags.DEFINE_float(‘learning_rate’, 1e-3, ‘learning rate’)
flags.DEFINE_integer(‘num_classes’, 20, ‘number of classes in the model’)
flags.DEFINE_boolean(‘multi_gpu’, False, ‘Use if wishing to train with more than 1 GPU.’)
flags.DEFINE_float(‘label_smoothing’, 0.02, ‘label smoothing’)
flags.DEFINE_integer(‘yolo_max_boxes’, 100, ‘yolo max boxes’)

def transform(image, label):label_encoder = anchor_label.encode(label)return image, label_encoder

def main(_argv):train_dataset = load_tfrecord_dataset(FLAGS.batch_size,FLAGS.dataset, FLAGS.classes, FLAGS.size)Yolo = Model(cfg=FLAGS.yaml_dir)anchors = Yolo.model[-1].anchorsstride = Yolo.model[-1].stridenum_classes = FLAGS.num_classes
anchor_label = AnchorLabeler(anchors,grids=FLAGS.img_size / stride,img_size=FLAGS.img_size,assign_method=‘wh’,extend_offset=‘True’)train_dataset = train_dataset.map(transform, num_parallel_calls=tf.data.experimental.AUTOTUNE)train_dataset = train_dataset.batch(FLAGS.batch_size).prefetch(tf.data.experimental.AUTOTUNE)print(train_dataset)Yolo_loss = Loss(anchors, iou_thres=0.3,num_classes=num_classes,label_smoothing=FLAGS.label_smoothing,img_size=FLAGS.img_size)optimizer = tf.keras.optimizers.Adam(lr=FLAGS.learning_rate)Yolo = Yolo(FLAGS.img_size)for epoch in range(0, FLAGS.epochs):for step, (image, target) in enumerate(train_dataset):with tf.GradientTape() as tape:output = Yolo(image) iou_loss, conf_loss, prob_loss = Yolo_loss(target, output)pred_loss = iou_loss+conf_loss+prob_losstotal_loss = tf.reduce_sum(pred_loss)grads = tape.gradient(total_loss, Yolo.trainable_variables)optimizer.apply_gradients(zip(grads, Yolo.trainable_variables))
logging.info("{}train{}, {}, {}".format(epoch, step, total_loss.numpy(),list(map(lambda x: np.sum(x.numpy()), pred_loss))))tf.saved_model.save(Yolo, ‘/data/Yolov5/weights/’)

if name==‘main’:app.run(main)

运行 train.py 脚本文件开始训练，此处要注意 cuda 和 cudnn的安装，使得 TensorFlow 能够成功调用 GPU 进行训练，效果如下。

训练完成后保存的网络模型 saved_model.pb 和 Variables 参数文件夹将保存在项目中的 weights 文件路径下。

3. 部

3.1 安装 OpenVINO™ 工具套件

登 录 https://www.intel.cn/content/www/cn/zh/developer/tools/openvino-toolkit/overview.html 官 网， 选 择 部 署 的操作系统和版本等进行下载和安装，本文的所有实现基于Windows 操作系统下的 2021.4.1 LTS 版本。

3.2 转换 OpenVINO™ 工具套件的 IR 格式

在 安 装 OpenVINO™ 工 具 套 件 后 的 目 录 下 \deployment_tools\model_optimizer 中， 通 过 mo_tf.py 脚本文件进行转换，由于此处TensorFlow保存得到的.pb只保存了网络结构，网络参数保存在 Variables 文件夹下，所以不能只将 .pb 文件作为输入进行转换，需要以 .pb 文件和Variables 文件夹的目录作为输入，并指定输出路径。如果 .pb没有明确输入图像的尺寸大小，需要加入 --input_shape 参数指定图像尺寸 ,–data_type 为设置数据格式为 32 位浮点数。

$python mo_tf.py --saved_model_dir <.pb文件夹路径>–input_shape [1,640,640,3] --output_dir <输出文件夹路径> --data_type FP32

运行成功之后会在输出文件夹路径下获得.xml和.bin文件，.xml 和 .bin 是 OpenVINO™ 工具套件中的模型存储方式，后续将基于 .bin 和 .xml 文件进行部署，效果如下。

3.3 推理部署

此实例将在 C++ 上进行推理部署， 在部署中包括有引擎初始化、数据准备、推理、结果处理等方面。引擎初始化需要读入转化后的模型文件并获取图像的输入输出信息。在数据准备中需要将输入图像缩放到 640*640 的尺寸大小并将通道输入改为 RGB。然后将输入填充在 blob 中，进行推理。得到 3 个检测头，分别对应 80、40 和 20 的栅格尺寸，并依次对结果进行解析。最后通过 NMS 剔除多余的候选框。

// 导入头文件
#include
#include
using namespace InferenceEngine;
using namespace std;
using namespace cv;
int main() {
// 初始化推理引擎
Core ie;
// 读取转换得到的 .xml 和 .bin 文件
CNNNetwork network = ie.ReadNetwork("./openvino/yolov5s.xml", “./openvino/yolov5s.bin”);

// 获取设置输入输出格式
// 从模型中获取输入数据的格式信息
InputsDataMap inputsInfo = network.getInputsInfo();InputInfo::Ptr& input = inputsInfo.begin()->second;string inputs_name = inputsInfo.begin()->first;
ICNNNetwork::InputShapes inputShapes = network.getInputShapes();network.reshape(inputShapes);

// 从模型中获取推断结果的格式
OutputsDataMap outputsInfo = network.getOutputsInfo();vector OutputsBlobs_names;for (auto& item_out : outputsInfo) {OutputsBlobs_names.push_back(item_out.first);item_out.second->setPrecision(Precision::FP32);}

// 获取可执行网络 , 这里的 CPU 指的是推断运行的器件 , 可选"GPU"ExecutableNetwork executable_network = ie.LoadNetwork(network, “CPU”);

// 推理请求
InferRequest infer_request = executable_network.CreateInferRequest();InferenceEngine::Blob::Ptr lrInputBlob = infer_request.GetBlob(inputs_name);float* buffer = lrInputBlob->buffer().as();
// 输入推理图像
Mat src = cv::imread("./img/test.jpg");size_t h = lrInputBlob->getTensorDesc().getDims()[2];size_t w = lrInputBlob->getTensorDesc().getDims()[3];size_t image_size = h * w;

Mat inframe = src.clone();cv::resize(src, src, Size(640, 640));cv::cvtColor(src, src, COLOR_BGR2RGB);InferenceEngine::LockedMemory blobMapped = InferenceEngine::asInferenceEngine::MemoryBlob(lrInputBlob)->wmap();float* blob_data = blobMapped.as();

//nchw
for (size_t row = 0; row < h; row++) {for (size_t col = 0; col < w; col++) {for (size_t ch = 0; ch < 3; ch++) {blob_data[image_size*ch + row * w + col] = float(src.at(row, col)[ch]) / 255.0f;
}
}
}

// 执行推理
infer_request.Infer();
// 设置置信度阈值和 NMS 阈值
float _cof_threshold = 0.1;float _nms_area_threshold = 0.5;
// 获取各层结果
vector origin_rect;
vector origin_rect_cof;int s[3] = { 80,40,20 };vector Blob::Ptr blobs;int i = 0;for (auto OutputsBlob_name : OutputsBlobs_names) {Blob::Ptr OutputBlob = infer_request.GetBlob(OutputsBlob_name);parse_yolov5(OutputBlob, s[i], _cof_threshold, origin_rect, origin_rect_cof);++i;}
// 后处理获得最终检测结果
vector final_id;

// 进行 NMS 处理，过滤重叠多余的候选框dnn::NMSBoxes(origin_rect, origin_rect_cof, _cof_threshold, _nms_area_threshold,
// 根据 final_id 获取最终结果
for (size_t i = 0; i < final_id.size(); ++i)
{
int idx = final_id[i];Rect box = origin_rect[idx];cv::rectangle(inframe, box, Scalar(140, 199, 0), 1, 8, 0);
}

cv::imwrite("./img/output.jpg", inframe);
}

此处要注意的是网络输出的结果需要进行转换处理，将中心点坐标转化为角点坐标和剔除置信度较低的候选框。

bool Detector::parse_yolov5(const Blob::Ptr &blob,int net_grid,float cof_threshold,vector & o_rect,vector& o_rect_cof){vector anchors = get_anchors(net_grid);

LockedMemory blobMapped = as(blob)->rmap();const float *output_blob = blobMapped.as();

//n 个类是 n+5
int item_size = 25;
size_t anchor_n = 3;
for(int n=0;n for(int i=0;i for(int j=0;j {
double box_prob = output_blob[nnet_gridnet_griditem_size + inet_grid*item_size + j *item_size+ 4];box_prob = sigmoid(box_prob);

// 框置信度不满足则整体置信度不满足
if(box_prob < cof_threshold)continue;

// 将中心点坐标转化为角点坐标
double x = output_blob[nnet_gridnet_grid*item_size +

inet_griditem_size + jitem_size + 0];double y = output_blob[nnet_gridnet_griditem_size +

inet_griditem_size + jitem_size + 1];double w = output_blob[nnet_gridnet_griditem_size +

inet_griditem_size + jitem_size + 2];double h = output_blob[nnet_gridnet_griditem_size +

inet_griditem_size + j item_size+ 3];double max_prob = 0;int idx=0;for(int t=5;t<25;++t){double tp= output_blob[nnet_gridnet_griditem_size +

inet_griditem_size + j *item_size+ t];
tp = sigmoid(tp);
if(tp > max_prob){
max_prob = tp;
idx = t;
}
}
float cof = box_prob * max_prob;
// 剔除边框置信度小于阈值的边框
if(cof < cof_threshold)
continue;

x = (sigmoid(x)*2 - 0.5 + j)*640.0f/net_grid;
y = (sigmoid(y)*2 - 0.5 + i)*640.0f/net_grid;
w = pow(sigmoid(w)2,2) * anchors[n2];
h = pow(sigmoid(h)2,2) * anchors[n2 + 1];
double r_x = x - w/2;
double r_y = y - h/2;Rect rect = Rect(round(r_x),round(r_y),round(w),round(h));
o_rect.push_back(rect);
o_rect_cof.push_back(cof);}
if(o_rect.size() == 0) return false;else return true;
}

double Detector::sigmoid(double x){return (1 / (1 + exp(-x)));
}

vector Detector::get_anchors(int net_grid){vector anchors(6);int anchor_80[6] = {10,13, 16,30, 33,23};int anchor_40[6] = {30,61, 62,45, 59,119};int anchor_20[6] = {116,90, 156,198, 373,326};if(net_grid == 80){ anchors.insert(anchors.begin(), anchor_80, anchor_80 + 6); }else if(net_grid == 40){ anchors.insert(anchors.begin(), anchor_40, anchor_40 + 6); }else if(net_grid == 20){ anchors.insert(anchors.begin(), anchor_20, anchor_20 + 6); }return anchors;}

运行以上的 main.cpp 工程，可以输出得到图像的检测结果画出候选的边界框并将结果保存为 img 文件夹下的 ouput.jpg。

4.性能测

Intel® DevCloud for the Edge 支持在英特尔的硬件平台上主动构建原型并试验面向计算机视觉的 AI 工作负载。可以使用 OpenVINO™ 工具套件以及 CPU、GPU 和 VPU 和 FPGA 的组合来测试模型的性能。Intel® DevCloud 使用 Jupyter* Notebook 直接在 web 浏览器中执行代码，并立即看到可视化结果。通过转换得到的 .xml 和 .bin 文件在不同边缘节点进行测试来分析性能，具体操作可以参考 https://bizwebcast.intel.cn/dev/articleDetails.html?id=95，测试结果见表4-1。