参考:每月深度1-2:YOLO V3-网络结构输入输出解析-史上最全
[代码剖析] 推荐阅读!
之前看了一遍 YOLO V3 的论文,写的挺有意思的
,尴尬的是,我这鱼的记忆,看完就忘了 
于是只能借助于代码,再看一遍细节了。
源码目录总览
tensorflow-yolov3-master
├── checkpoint //保存模型的目录
├── convert_weight.py//对权重进行转换,为了模型的预训练
├── core//核心代码文件夹
│ ├── backbone.py
│ ├── common.py
│ ├── config.py//配置文件
│ ├── dataset.py//数据处理
│ ├── __init__.py
│ ├── __pycache__
│ │ ├── backbone.cpython-36.pyc
│ │ ├── common.cpython-36.pyc
│ │ ├── config.cpython-36.pyc
│ │ ├── dataset.cpython-36.pyc
│ │ ├── __init__.cpython-36.pyc
│ │ ├── utils.cpython-36.pyc
│ │ └── yolov3.cpython-36.pyc
│ ├── utils.py
│ └── yolov3.py//网络核心结构
├── data
│ ├── anchors//预训练框
│ │ ├── basline_anchors.txt
│ │ └── coco_anchors.txt
│ ├── classes//训练预测目标的种类
│ │ ├── coco.names
│ │ └── voc.names
│ ├── dataset//保存图片的相关信息:路径,box,置信度,类别编号
│ │ ├── voc_test.txt//测试数据
│ │ └── voc_train.txt//训练数据
│ └── log//存储log
│ └── events.out.tfevents.1564706916.WIN-RCRPPSUQJFP
├── docs//比较混杂
│ ├── Box-Clustering.ipynb//根据数据信息生成预选框anchors
│ ├── images
│ │ ├── 611_result.jpg
│ │ ├── darknet53.png
│ │ ├── iou.png
│ │ ├── K-means.png
│ │ ├── levio.jpeg
│ │ ├── probability_extraction.png
│ │ ├── road.jpeg
│ │ ├── road.mp4
│ │ └── yolov3.png
│ └── requirements.txt//环境搭建
├── evaluate.py//模型评估
├── freeze_graph.py//生成pb文件
├── image_demo.py//一张图片测试的demo
├── LICENSE
├── LICENSE.fuck
├── mAP//模型评估相关信息存储
│ ├── extra
│ │ ├── class_list.txt
│ │ ├── convert_gt_xml.py
│ │ ├── convert_gt_yolo.py
│ │ ├── convert_keras-yolo3.py
│ │ ├── convert_pred_darkflow_json.py
│ │ ├── convert_pred_yolo.py
│ │ ├── find_class.py
│ │ ├── intersect-gt-and-pred.py
│ │ ├── README.md
│ │ ├── remove_class.py
│ │ ├── remove_delimiter_char.py
│ │ ├── remove_space.py
│ │ ├── rename_class.py
│ │ └── result.txt
│ ├── __init__.py
│ └── main.py
├── README.md
├── scripts
│ ├── show_bboxes.py
│ └── voc_annotation.py//把xml转化为网络可以使用的txt文件
├── train.py//模型训练
└── video_demo.py//视屏测试的demo
接下来,我按照看代码的顺序来详细说明了。
core/dataset.py

#! /usr/bin/env python
# coding=utf-8
# ================================================================
# Copyright (C) 2019 * Ltd. All rights reserved.
#
# Editor : VIM
# File name : dataset.py
# Author : YunYang1994
# Created date: 2019-03-15 18:05:03
# Description :
#
# ================================================================
import os
import cv2
import random
import numpy as np
import tensorflow as tf
import core.utils as utils
from core.config import cfg
class DataSet(object):
"""implement Dataset here"""
def __init__(self, dataset_type):
self.annot_path = cfg.TRAIN.ANNOT_PATH if dataset_type == 'train' else cfg.TEST.ANNOT_PATH
self.input_sizes = cfg.TRAIN.INPUT_SIZE if dataset_type == 'train' else cfg.TEST.INPUT_SIZE
self.batch_size = cfg.TRAIN.BATCH_SIZE if dataset_type == 'train' else cfg.TEST.BATCH_SIZE
self.data_aug = cfg.TRAIN.DATA_AUG if dataset_type == 'train' else cfg.TEST.DATA_AUG
self.train_input_sizes = cfg.TRAIN.INPUT_SIZE
self.strides = np.array(cfg.YOLO.STRIDES)
self.classes = utils.read_class_names(cfg.YOLO.CLASSES)
self.num_classes = len(self.classes)
self.anchors = np.array(utils.get_anchors(cfg.YOLO.ANCHORS))
self.anchor_per_scale = cfg.YOLO.ANCHOR_PER_SCALE
self.max_bbox_per_scale = 150
self.annotations = self.load_annotations(dataset_type) # read and shuffle annotations
self.num_samples = len(self.annotations) # dataset size
self.num_batchs = int(np.ceil(self.num_samples / self.batch_size)) # 向上取整
self.batch_count = 0 # batch index
def load_annotations(self, dataset_type):
with open(self.annot_path, 'r') as f:
txt = f.readlines()
annotations = [line.strip() for line in txt if len(line.strip().split()[1:]) != 0]
# np.random.seed(1) # for debug
np.random.shuffle(annotations)
return annotations
def __iter__(self):
return self
def next(self):
with tf.device('/cpu:0'):
self.train_input_size_h, self.train_input_size_w = random.choice(self.train_input_sizes)
self.train_output_sizes_h = self.train_input_size_h // self.strides
self.train_output_sizes_w = self.train_input_size_w // self.strides
# ================================================================ #
batch_image = np.zeros((self.batch_size, self.train_input_size_h, self.train_input_size_w, 3))
batch_label_sbbox = np.zeros((self.batch_size, self.train_output_sizes_h[0], self.train_output_sizes_w[0],
self.anchor_per_scale, 5 + self.num_classes))
batch_label_mbbox = np.zeros((self.batch_size, self.train_output_sizes_h[1], self.train_output_sizes_w[1],
self.anchor_per_scale, 5 + self.num_classes))
batch_label_lbbox = np.zeros((self.batch_size, self.train_output_sizes_h[2], self.train_output_sizes_w[2],
self.anchor_per_scale, 5 + self.num_classes))
# ================================================================ #
batch_sbboxes = np.zeros((self.batch_size, self.max_bbox_per_scale, 4))
batch_mbboxes = np.zeros((self.batch_size, self.max_bbox_per_scale, 4))
batch_lbboxes = np.zeros((self.batch_size, self.max_bbox_per_scale, 4))
num = 0 # sample in one batch's index
if self.batch_count < self.num_batchs:
while num < self.batch_size:
index = self.batch_count * self.batch_size + num
if index >= self.num_samples: # 从头开始
index -= self.num_samples
annotation = self.annotations[index]
# 样本预处理
image, bboxes = self.parse_annotation(annotation)
# Anchor & GT 匹配
label_sbbox, label_mbbox, label_lbbox, sbboxes, mbboxes, lbboxes = self.preprocess_true_boxes(
bboxes)
batch_image[num, :, :, :] = image
batch_label_sbbox[num, :, :, :, :] = label_sbbox
batch_label_mbbox[num, :, :, :, :] = label_mbbox
batch_label_lbbox[num, :, :, :, :] = label_lbbox
batch_sbboxes[num, :, :] = sbboxes
batch_mbboxes[num, :, :] = mbboxes
batch_lbboxes[num, :, :] = lbboxes
num += 1
self.batch_count += 1
return batch_image, batch_label_sbbox, batch_label_mbbox, batch_label_lbbox, \
batch_sbboxes, batch_mbboxes, batch_lbboxes
else:
self.batch_count = 0
np.random.shuffle(self.annotations)
raise StopIteration
def random_horizontal_flip(self, image, bboxes):
if random.random() < 0.5:
_, w, _ = image.shape
image = image[:, ::-1, :]
bboxes[:, [0, 2]] = w - bboxes[:, [2, 0]]
return image, bboxes
def random_crop(self, image, bboxes):
if random.random() < 0.5:
h, w, _ = image.shape
max_bbox = np.concatenate([np.min(bboxes[:, 0:2], axis=0), np.max(bboxes[:, 2:4], axis=0)], axis=-1)
max_l_trans = max_bbox[0]
max_u_trans = max_bbox[1]
max_r_trans = w - max_bbox[2]
max_d_trans = h - max_bbox[3]
crop_xmin = max(0, int(max_bbox[0] - random.uniform(0, max_l_trans)))
crop_ymin = max(0, int(max_bbox[1] - random.uniform(0, max_u_trans)))
crop_xmax = max(w, int(max_bbox[2] + random.uniform(0, max_r_trans)))
crop_ymax = max(h, int(max_bbox[3] + random.uniform(0, max_d_trans)))
image = image[crop_ymin: crop_ymax, crop_xmin: crop_xmax]
bboxes[:, [0, 2]] = bboxes[:, [0, 2]] - crop_xmin
bboxes[:, [1, 3]] = bboxes[:, [1, 3]] - crop_ymin
return image, bboxes
def random_translate(self, image, bboxes):
if random.random() < 0.5:
h, w, _ = image.shape
max_bbox = np.concatenate([np.min(bboxes[:, 0:2], axis=0), np.max(bboxes[:, 2:4], axis=0)], axis=-1)
max_l_trans = max_bbox[0]
max_u_trans = max_bbox[1]
max_r_trans = w - max_bbox[2]
max_d_trans = h - max_bbox[3]
tx = random.uniform(-(max_l_trans - 1), (max_r_trans - 1))
ty = random.uniform(-(max_u_trans - 1), (max_d_trans - 1))
M = np.array([[1, 0, tx], [0, 1, ty]])
image = cv2.warpAffine(image, M, (w, h))
bboxes[:, [0, 2]] = bboxes[:, [0, 2]] + tx
bboxes[:, [1, 3]] = bboxes[:, [1, 3]] + ty
return image, bboxes
def parse_annotation(self, annotation):
line = annotation.split()
image_path = line[0]
if not os.path.exists(image_path):
raise KeyError("%s does not exist ... " % image_path)
image = np.array(cv2.imread(image_path))
bboxes = np.array([list(map(int, box.split(','))) for box in line[1:]])
if self.data_aug:
image, bboxes = self.random_horizontal_flip(np.copy(image), np.copy(bboxes))
image, bboxes = self.random_crop(np.copy(image), np.copy(bboxes))
image, bboxes = self.random_translate(np.copy(image), np.copy(bboxes))
image, bboxes = utils.image_preporcess(np.copy(image), [self.train_input_size_h, self.train_input_size_w],
np.copy(bboxes))
return image, bboxes
def bbox_iou(self, boxes1, boxes2):
boxes1 = np.array(boxes1)
boxes2 = np.array(boxes2)
boxes1_area = boxes1[..., 2] * boxes1[..., 3]
boxes2_area = boxes2[..., 2] * boxes2[..., 3]
boxes1 = np.concatenate([boxes1[..., :2] - boxes1[..., 2:] * 0.5,
boxes1[..., :2] + boxes1[..., 2:] * 0.5], axis=-1)
boxes2 = np.concatenate([boxes2[..., :2] - boxes2[..., 2:] * 0.5,
boxes2[..., :2] + boxes2[..., 2:] * 0.5], axis=-1)
left_up = np.maximum(boxes1[..., :2], boxes2[..., :2])
right_down = np.minimum(boxes1[..., 2:], boxes2[..., 2:])
inter_section = np.maximum(right_down - left_up, 0.0)
inter_area = inter_section[..., 0] * inter_section[..., 1]
union_area = boxes1_area + boxes2_area - inter_area
return inter_area / union_area
def preprocess_true_boxes(self, bboxes):
# ================================================================ #
label = [np.zeros((self.train_output_sizes_h[i], self.train_output_sizes_w[i], self.anchor_per_scale,
5 + self.num_classes)) for i in range(3)]
""" match info
hypothesis input size 320 x 480, label dim
| 40 x 60 x 3 x 17 |
| 20 x 30 x 3 x 17 |
| 10 x 15 x 3 x 17 |
"""
bboxes_xywh = [np.zeros((self.max_bbox_per_scale, 4)) for _ in range(3)]
""" match gt set
bboxes_xywh dim
| 3 x 150 x 4 |
"""
bbox_count = np.zeros((3,))
# ================================================================ #
for bbox in bboxes:
bbox_coor = bbox[:4] # xmin, ymin, xmax, ymax
bbox_class_ind = bbox[4] # class
# smooth onehot label
onehot = np.zeros(self.num_classes, dtype=np.float)
onehot[bbox_class_ind] = 1.0
uniform_distribution = np.full(self.num_classes, 1.0 / self.num_classes)
deta = 0.01
smooth_onehot = onehot * (1 - deta) + deta * uniform_distribution
# box transform into 3 feature maps [center_x, center_y, w, h]
bbox_xywh = np.concatenate([(bbox_coor[2:] + bbox_coor[:2]) * 0.5, bbox_coor[2:] - bbox_coor[:2]], axis=-1)
bbox_xywh_scaled = 1.0 * bbox_xywh[np.newaxis, :] / self.strides[:, np.newaxis]
# =========================== match iou ========================== #
iou = [] # 3x3
exist_positive = False
for i in range(3): # different feature map
anchors_xywh = np.zeros((self.anchor_per_scale, 4))
anchors_xywh[:, 0:2] = np.floor(bbox_xywh_scaled[i, 0:2]).astype(np.int32) + 0.5
anchors_xywh[:, 2:4] = self.anchors[i]
iou_scale = self.bbox_iou(bbox_xywh_scaled[i][np.newaxis, :], anchors_xywh)
iou.append(iou_scale)
iou_mask = iou_scale > 0.3
if np.any(iou_mask):
xind, yind = np.floor(bbox_xywh_scaled[i, 0:2]).astype(np.int32)
label[i][yind, xind, iou_mask, :] = 0
label[i][yind, xind, iou_mask, 0:4] = bbox_xywh
label[i][yind, xind, iou_mask, 4:5] = 1.0
label[i][yind, xind, iou_mask, 5:] = smooth_onehot
bbox_ind = int(bbox_count[i] % self.max_bbox_per_scale)
bboxes_xywh[i][bbox_ind, :4] = bbox_xywh
bbox_count[i] += 1
exist_positive = True
if not exist_positive:
best_anchor_ind = np.argmax(np.array(iou).reshape(-1), axis=-1)
best_detect = int(float(best_anchor_ind) / self.anchor_per_scale)
best_anchor = int(best_anchor_ind % self.anchor_per_scale)
xind, yind = np.floor(bbox_xywh_scaled[best_detect, 0:2]).astype(np.int32)
label[best_detect][yind, xind, best_anchor, :] = 0
label[best_detect][yind, xind, best_anchor, 0:4] = bbox_xywh
label[best_detect][yind, xind, best_anchor, 4:5] = 1.0
label[best_detect][yind, xind, best_anchor, 5:] = smooth_onehot
bbox_ind = int(bbox_count[best_detect] % self.max_bbox_per_scale)
bboxes_xywh[best_detect][bbox_ind, :4] = bbox_xywh
bbox_count[best_detect] += 1
label_sbbox, label_mbbox, label_lbbox = label # different size feature map's anchor match info
sbboxes, mbboxes, lbboxes = bboxes_xywh # different size feature map's matched gt set
return label_sbbox, label_mbbox, label_lbbox, sbboxes, mbboxes, lbboxes
def __len__(self):
return self.num_batchs
if __name__ == '__main__':
val = DataSet('test')
for idx in range(val.num_batchs):
batch_image, batch_label_sbbox, batch_label_mbbox, batch_label_lbbox, \
batch_sbboxes, batch_mbboxes, batch_lbboxes = val.next()
print('# ================================================================ #')
这部分是用来加载数据的。
整个 core/dataset.py 实现了一个 DataSet 类,每个 batch 的数据通过迭代函数 next() 获得。函数返回
batch_image # one batch resized images batch_label_sbbox # 第一个尺度下的匹配结果 batch_label_mbbox # 第二个尺度下的匹配结果 batch_label_lbbox # 第三个尺度下的匹配结果 batch_sbboxes # 第一个尺度下匹配的 GT 集合 batch_mbboxes # 第二个尺度下匹配的 GT 集合 batch_lbboxes # 第三个尺度下匹配的 GT 集合
其中 batch_image 为网络输入特征,按照 NxHxWxC 维度排列,batch_label_sbbox、batch_label_mbbox 、batch_label_lbbox 这三个用以确定不同尺度下的 anchor 做正样本还是负样本;batch_sbboxes、batch_mbboxes、batch_lbboxes 这三个就有点意思了,是为了后续区分负样本是否有可能变成正样本(回归到 bbox 了) 做准备。
,可以继续往下看。数据集在 DataSet 初始化的时候就被打乱了,如果你想每次调试运行看看上面的代码是怎么工作的,可以再 load_annotations() 函数里固定随机因子
np.random.seed(1)
数据格式
不管你训练什么数据集(VOC、COCO),数据的标注格式都要转换成以下格式 ' image_path [bbox class] ...':
D:/tyang/drive0703/JPEGImages/test_dataset_without2018/2017_07_24_10_40_019934.jpg 658,335,970,601,0 350,327,577,480,0 137,311,378,453,0 526,316,611,430,3 1334,485,1544,634,5 1089,207,1119,324,8 808,132,836,193,8 688,121,713,182,8 137,156,182,236,8 56,270,93,389,7 100,263,139,397,7 1097,357,1142,517,7 1275,376,1353,558,7 1391,350,1444,484,7 1379,392,1493,621,7 1457,425,1571,602,7 1558,363,1591,483,7 1575,392,1643,566,7 1829,373,1899,533,7
当然,scripts/voc_annotation.py 提供了将 VOC 的 xml 格式转换成需求格式,这部分不是很难,其他格式的数据集,自己瞎写写转一转也没啥问题。
数据预处理
对于每个样本,先经过预处理函数 parse_annotation() 加载 image & bboxes,对于训练数据,这里会执行一些 data argument。
Anchor & GT 匹配
然后,最重要的部分来了,通过 preprocess_true_boxes() 来实现 Anchor & GT 匹配:
label = [np.zeros((self.train_output_sizes_h[i], self.train_output_sizes_w[i], self.anchor_per_scale,
5 + self.num_classes)) for i in range(3)]
""" match info
hypothesis input size 320 x 480, label dim
| 40 x 60 x 3 x 17 |
| 20 x 30 x 3 x 17 |
| 10 x 15 x 3 x 17 |
"""
bboxes_xywh = [np.zeros((self.max_bbox_per_scale, 4)) for _ in range(3)]
""" match gt set
bboxes_xywh dim
| 3 x 150 x 4 |
"""
label 里保存的是三个尺度下的 Anchor 的匹配情况;bboxes_xywh 里则保存的是三个尺度小被匹配上的 GT 集合。
网络在三个尺度的 feature map 下检测目标,这三个尺度的 feature map 由输入图片分别经过 stride=8, 16, 32 获得。
假定训练中输入尺度是 320x480,那么这三个 feature map size: 40x60, 20x30, 10x15。
对于每个 feature map 下,作者分别设计(聚类)了三种 anchor,例如 data/anchors/basline_anchors.txt 中:
# small 1.25,1.625 2.0,3.75 4.125,2.875 # middle 1.875,3.8125 3.875,2.8125 3.6875,7.4375 # large 3.625,2.8125 4.875,6.1875 11.65625,10.1875
因此,第一个尺度里有 40 x 60 x 3 = 7200 个 anchor;第二个尺度里有 20 x 30 x 3 = 1800 个 anchor;第三个尺度里有 10x 15x 3 = 450 个 anchor。
值得注意的是,原始 bbox 是按照 [xmin, ymin, xmax, ymax],需要转换成 [center_x, center_y, w, h],即:
[164 99 242 178] -> [203. 138.5 78. 79.] [ 87 96 144 142] -> [115.5 119. 57. 46.] [ 34 92 94 134] -> [ 64. 113. 60. 42.] [131 93 152 127] -> [141.5 110. 21. 34.] [333 143 386 187] -> [359.5 165. 53. 44.] [272 61 279 96] -> [275.5 78.5 7. 35.] [202 39 209 57] -> [205.5 48. 7. 18.] [172 35 178 53] -> [175. 44. 6. 18.] [ 34 46 45 69] -> [39.5 57.5 11. 23.] [ 14 80 23 115] -> [18.5 97.5 9. 35.] [ 25 77 34 117] -> [29.5 97. 9. 40.] [274 105 285 153] -> [279.5 129. 11. 48.] [318 111 338 165] -> [328. 138. 20. 54.] [347 103 361 143] -> [354. 123. 14. 40.] [344 116 373 184] -> [358.5 150. 29. 68.] [364 125 392 178] -> [378. 151.5 28. 53.] [389 107 397 143] -> [393. 125. 8. 36.] [393 116 410 167] -> [401.5 141.5 17. 51.] [457 110 474 157] -> [465.5 133.5 17. 47.]
label 最后一个维度是 5 + self.num_classes,其中 5, 前 4 维是 [center_x, center_y, w, h] GT bbox, 第 5 维是 0/1, 0 表示无匹配,1 表示匹配成功。self.num_classes 用来表示目标类别,之所以要用这么多维数据来表示,是因为将整形 label 转换成了 one-hot 形式。同时这里做了 label smooth 操作,例如 label 0 ->[9.90833333e-01 8.33333333e-04 8.33333333e-04 8.33333333e-04 8.33333333e-04 8.33333333e-04 8.33333333e-04 8.33333333e-04 8.33333333e-04 8.33333333e-04 8.33333333e-04 8.33333333e-04]。
考虑到在每个尺度上,每个 GT bbox 最多只能有一个 match anchor,因此 YOLO V3 的匹配过程和 SSD 的匹配过程有所差异,对于每个 bbox(下面会以 [203. 138.5 78. 79.] 这个 bbox 为例):
- 将 bbox 映射到每个 feature map 上,获得 bbox_xywh_scaled :
[[25.375 17.3125 9.75 9.875 ] [12.6875 8.65625 4.875 4.9375 ] [ 6.34375 4.328125 2.4375 2.46875 ]]
- 在每个尺度上尝试匹配,只利用中心的在 bbox 中心点最近的 anchor 尝试匹配,例如第二个尺度上的 bbox=[12.6875 8.65625 4.875 4.9375] 将尝试 anchor 集合(anchors_xywh) 进行匹配:
[[12.5 8.5 1.875 3.8125] [12.5 8.5 3.875 2.8125] [12.5 8.5 3.6875 7.4375]]
- 匹配计算 bbox_iou,如果找到满足大于 0.3 的一对,即为匹配成功 (False, True, True)。匹配成功后就往 label 和 bboxes_xywh 里填信息就好了。
""" label[1][8, 12, [False True True], :] = 0 label[1][8, 12, [False True True], 0:4] = [203. 138.5 78. 79.] label[1][8, 12, [False True True], 4:5] = 1.0 label[1][8, 12, [False True True], 5:] = [9.90833333e-01 8.33333333e-04 8.33333333e-04 8.33333333e-04 8.33333333e-04 8.33333333e-04 8.33333333e-04 8.33333333e-04 8.33333333e-04 8.33333333e-04 8.33333333e-04 8.33333333e-04] bboxes_xywh[1][bbox_ind, :4] = [203. 138.5 78. 79. ] """ bbox_ind 是这个尺度下累计匹配成功的 GT 数量,代码中限制范围 [0-149] - 如果有 bbox 在各个 feature map 都找不到满足的匹配 anchor,那就退而求其次,在所有 feature map 的 anchor 里寻找一个最大匹配就好了。
到此,匹配结束。
train.py
INPUT
输入层对应 dataset.py 每个batch 返回的变量
with tf.name_scope('define_input'):
self.input_data = tf.placeholder(dtype=tf.float32, name='input_data', shape=[None, None, None, 3])
self.label_sbbox = tf.placeholder(dtype=tf.float32, name='label_sbbox')
self.label_mbbox = tf.placeholder(dtype=tf.float32, name='label_mbbox')
self.label_lbbox = tf.placeholder(dtype=tf.float32, name='label_lbbox')
self.true_sbboxes = tf.placeholder(dtype=tf.float32, name='sbboxes')
self.true_mbboxes = tf.placeholder(dtype=tf.float32, name='mbboxes')
self.true_lbboxes = tf.placeholder(dtype=tf.float32, name='lbboxes')
self.trainable = tf.placeholder(dtype=tf.bool, name='training')
MODEL & LOSS
YOLOV3 的 loss 分为三部分,回归 loss, 二分类(前景/背景) loss, 类别分类 loss
with tf.name_scope("define_loss"):
self.model = YOLOV3(self.input_data, self.trainable, self.net_flag)
self.net_var = tf.global_variables()
self.giou_loss, self.conf_loss, self.prob_loss = self.model.compute_loss(self.label_sbbox,
self.label_mbbox,
self.label_lbbox,
self.true_sbboxes,
self.true_mbboxes,
self.true_lbboxes)
self.loss = self.giou_loss + self.conf_loss + self.prob_loss
Learning rate
with tf.name_scope('learn_rate'):
self.global_step = tf.Variable(1.0, dtype=tf.float64, trainable=False, name='global_step')
warmup_steps = tf.constant(self.warmup_periods * self.steps_per_period,
dtype=tf.float64, name='warmup_steps') # warmup_periods epochs
train_steps = tf.constant((self.first_stage_epochs + self.second_stage_epochs) * self.steps_per_period,
dtype=tf.float64, name='train_steps')
self.learn_rate = tf.cond(
pred=self.global_step < warmup_steps,
true_fn=lambda: self.global_step / warmup_steps * self.learn_rate_init,
false_fn=lambda: self.learn_rate_end + 0.5 * (self.learn_rate_init - self.learn_rate_end) * (
1 + tf.cos((self.global_step - warmup_steps) / (train_steps - warmup_steps) * np.pi)))
global_step_update = tf.assign_add(self.global_step, 1.0)
"""
训练分为两个阶段,第一阶段里前面又划分出一段作为“热身阶段”:
热身阶段:learn_rate = (global_step / warmup_steps) * learn_rate_init
其他阶段:learn_rate_end + 0.5 * (learn_rate_init - learn_rate_end) * (
1 + tf.cos((global_step - warmup_steps) / (train_steps - warmup_steps) * np.pi))
"""
假定遍历一遍数据集需要 100batch, warmup_periods=2, first_stage_epochs=20, second_stage_epochs=30, learn_rate_init=1e-4, learn_rate_end=1e-6, 那么整个训练过程中学习率是这样的:

import numpy as np
import matplotlib.pyplot as plt
steps_per_period = 100
warmup_periods=2
first_stage_epochs=20
second_stage_epochs=30
learn_rate_init=1e-4
learn_rate_end=1e-6
warmup_steps = warmup_periods * steps_per_period
train_steps = (first_stage_epochs + second_stage_epochs) * steps_per_period
def learn_rate_strategy(global_step, warmup_steps, train_steps,
learn_rate_init, learn_rate_end):
"""
:param global_step:
:param warmup_steps:
:param learn_rate_init:
:param learn_rate_end:
:return:
"""
if global_step < warmup_steps:
learn_rate = (global_step / warmup_steps) * learn_rate_init
else:
learn_rate = learn_rate_end + 0.5 * (learn_rate_init - learn_rate_end) * (
1 + np.cos((global_step - warmup_steps) / (train_steps - warmup_steps) * np.pi))
return learn_rate
learn_rate_list = []
for step in range(train_steps):
learing_rate = learn_rate_strategy(step, warmup_steps, train_steps,
learn_rate_init, learn_rate_end)
learn_rate_list.append(learing_rate)
step = range(train_steps)
print(learn_rate_list[-1])
plt.plot(step, learn_rate_list, 'g-', linewidth=2, label='learing_rate')
plt.xlabel('step')
plt.ylabel('learing rate')
plt.legend(loc='upper right')
plt.tight_layout()
plt.show()

two stage train
整个训练按照任务划分成了两个阶段,之所以这么设计,是考虑作者是拿原始的 DarkNet 来 finetune 的。
finetune 的一般流程就是,利用预训练的模型赋初值,先固定 backbone,只训练最后的分类/回归层。然后放开全部训练。
也可以对于浅层特征可以用小的学习率来微调(因为网络里浅层特征提取的边界纹理信息可能都是相近的,不需要作大调整),越接近于输出层可能需要调整的越多,输出层因为没有用其他模型初始化(随机初始化),因此需要从头训练。
for epoch in range(1, 1 + self.first_stage_epochs + self.second_stage_epochs):
if epoch <= self.first_stage_epochs:
train_op = self.train_op_with_frozen_variables
else:
train_op = self.train_op_with_all_variables
first_stage_train
这个阶段将专注于训练最后的检测部分,即分类和回归
with tf.name_scope("define_first_stage_train"):
self.first_stage_trainable_var_list = []
for var in tf.trainable_variables():
var_name = var.op.name
var_name_mess = str(var_name).split('/')
if var_name_mess[0] in ['conv_sbbox', 'conv_mbbox', 'conv_lbbox']:
self.first_stage_trainable_var_list.append(var)
first_stage_optimizer = tf.train.AdamOptimizer(self.learn_rate).minimize(self.loss,
var_list=self.first_stage_trainable_var_list)
with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
with tf.control_dependencies([first_stage_optimizer, global_step_update]):
with tf.control_dependencies([moving_ave]):
self.train_op_with_frozen_variables = tf.no_op()
second_stage_train
这个阶段就是整体训练,没什么好说的
with tf.name_scope("define_second_stage_train"):
second_stage_trainable_var_list = tf.trainable_variables()
second_stage_optimizer = tf.train.AdamOptimizer(self.learn_rate).minimize(self.loss,
var_list=second_stage_trainable_var_list)
with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
with tf.control_dependencies([second_stage_optimizer, global_step_update]):
with tf.control_dependencies([moving_ave]):
self.train_op_with_all_variables = tf.no_op()
ExponentialMovingAverage
with tf.name_scope("define_weight_decay"):
moving_ave = tf.train.ExponentialMovingAverage(self.moving_ave_decay).apply(tf.trainable_variables())
这个我涉世未深,还不甚明了
,参见 tf.train.ExponentialMovingAverage
core/backbone.py
这里定义了 Darknet-53 的主题框架
| 网络结构 | 代码结构 |
|---|---|
![]() |
![]() |
当然,你也可以根据你的需要,定义一些其他的 backbone,例如 mobilenet_v1、mobilenet_v2。
core/yolov3.py
这里是整个 YOLOV3 代码的灵魂之处了。
__build_nework
YOLOV3 同 SSD 一样,是多尺度目标检测。选择了 stride=8, 16, 32 三个尺度的 feature map 来设计 anchor, 以便分别实现对大、中和小物体的预测。
假定输入尺度是 320x480,那么这三个 feature map 的大小就是 40x60, 20x30, 10x15。分别在这三个尺度的 feature map 的基础上,通过 3*(4+1+classes) 个 3x3 的卷积核卷积来预测分类和回归结果。这里 3 代表每个尺度下设计了 3 中不同尺寸的 anchor,4 是 bbox 回归预测,1 则是表示该 anchor 是否包含目标,classes 则是你的数据集里具体的类别数量了。如此,可以推测出每个尺度下的预测输出维度为(我的数据集包含 12 个类别目标):
batch_size x 40 x 60 x 51 batch_size x 20 x 30 x 51 batch_size x 10 x 15 x 51
这些预测输出将和 core/dataset.py 文件里获得的 GT 信息作比较,计算 loss。
upsample
网络在 backbone 特征提取的基础上加上了上采样特征连接,加强了浅层特征表示。
def upsample(input_data, name, method="deconv"):
assert method in ["resize", "deconv"]
if method == "resize":
with tf.variable_scope(name):
input_shape = tf.shape(input_data)
output = tf.image.resize_nearest_neighbor(input_data, (input_shape[1] * 2, input_shape[2] * 2))
if method == "deconv":
# replace resize_nearest_neighbor with conv2d_transpose To support TensorRT optimization
numm_filter = input_data.shape.as_list()[-1]
output = tf.layers.conv2d_transpose(input_data, numm_filter, kernel_size=2, padding='same',
strides=(2, 2), kernel_initializer=tf.random_normal_initializer())
这里提供了两种实现方式,最近邻缩放和反卷积。
decode
同 SSD 类似,anchor 的回归并非直接坐标回归,而是通过编码后进行回归:

我们知道,检测框实际上是在先验框的基础上回归出来的。如上图所示:在其中一个输出尺度下的 feature map 上,有一个黑色的先验框($c_x, c_y, p_w, p_h$),其中 $c_x$ 和 $c_y$ 分别表示中心网格距离图像左上角的距离,$p_w$ 和 $p_h$ 则分别表示先验框的宽和高。
记网络回归输出为($t_x, t_y, t_w, t_h$),其中$t_x$ 和 $t_y$ 用以偏移先验框的中心到检测框,$p_w$ 和 $p_h$ 则用来缩放先验框到检测框大小,那么蓝色的检测框($b_x, b_y, b_w, b_h$)可以用以下表达式表示:
\begin{equation}
\label{a}
\begin{split}
& b_x = \sigma(t_x) + c_x \\
& b_y = \sigma(t_y) + c_y \\
& b_w = p_w e^{t_w} \\
& b_h = p_h e^{t_h} \\
\end{split}
\end{equation}
具体实现:
def decode(self, conv_output, anchors, stride):
"""
return tensor of shape [batch_size, output_size, output_size, anchor_per_scale, 5 + num_classes]
contains (x, y, w, h, score, probability)
"""
conv_shape = tf.shape(conv_output)
batch_size = conv_shape[0]
output_size_h = conv_shape[1]
output_size_w = conv_shape[2]
anchor_per_scale = len(anchors)
conv_output = tf.reshape(conv_output,
(batch_size, output_size_h, output_size_w, anchor_per_scale, 5 + self.num_class))
conv_raw_dxdy = conv_output[:, :, :, :, 0:2]
conv_raw_dwdh = conv_output[:, :, :, :, 2:4]
conv_raw_conf = conv_output[:, :, :, :, 4:5]
conv_raw_prob = conv_output[:, :, :, :, 5:]
# 划分网格
y = tf.tile(tf.range(output_size_h, dtype=tf.int32)[:, tf.newaxis], [1, output_size_w])
x = tf.tile(tf.range(output_size_w, dtype=tf.int32)[tf.newaxis, :], [output_size_h, 1])
xy_grid = tf.concat([x[:, :, tf.newaxis], y[:, :, tf.newaxis]], axis=-1)
xy_grid = tf.tile(xy_grid[tf.newaxis, :, :, tf.newaxis, :], [batch_size, 1, 1, anchor_per_scale, 1])
xy_grid = tf.cast(xy_grid, tf.float32) # 计算网格左上角的位置
# 根据论文公式计算预测框的中心位置
pred_xy = (tf.sigmoid(conv_raw_dxdy) + xy_grid) * stride
# 根据论文公式计算预测框的长和宽大小
pred_wh = (tf.exp(conv_raw_dwdh) * anchors) * stride
# 合并边界框的位置和长宽信息
pred_xywh = tf.concat([pred_xy, pred_wh], axis=-1)
pred_conf = tf.sigmoid(conv_raw_conf) # 计算预测框里object的置信度
pred_prob = tf.sigmoid(conv_raw_prob) # 计算预测框里object的类别概率
return tf.concat([pred_xywh, pred_conf, pred_prob], axis=-1)
compute_loss & loss_layer
代码分别从三个尺度出发,分别计算 边界框损失(giou_loss)、是否包含目标的置信度损失(conf_loss)以及具体类别的分类损失(prob_loss)
def compute_loss(self, label_sbbox, label_mbbox, label_lbbox, true_sbbox, true_mbbox, true_lbbox):
with tf.name_scope('smaller_box_loss'):
loss_sbbox = self.loss_layer(self.conv_sbbox, self.pred_sbbox, label_sbbox, true_sbbox,
anchors=self.anchors[0], stride=self.strides[0])
with tf.name_scope('medium_box_loss'):
loss_mbbox = self.loss_layer(self.conv_mbbox, self.pred_mbbox, label_mbbox, true_mbbox,
anchors=self.anchors[1], stride=self.strides[1])
with tf.name_scope('bigger_box_loss'):
loss_lbbox = self.loss_layer(self.conv_lbbox, self.pred_lbbox, label_lbbox, true_lbbox,
anchors=self.anchors[2], stride=self.strides[2])
with tf.name_scope('giou_loss'):
giou_loss = loss_sbbox[0] + loss_mbbox[0] + loss_lbbox[0]
with tf.name_scope('conf_loss'):
conf_loss = loss_sbbox[1] + loss_mbbox[1] + loss_lbbox[1]
with tf.name_scope('prob_loss'):
prob_loss = loss_sbbox[2] + loss_mbbox[2] + loss_lbbox[2]
return giou_loss, conf_loss, prob_loss
GIoU Loss
同 SSD(smooth L1 loss) 等检测算法相比,这里使用 GIoU 来衡量检测框和 GT bbox 之间的差距,具体可以参考论文和本代码作者的解读。

giou = tf.expand_dims(self.bbox_giou(pred_xywh, label_xywh), axis=-1)
input_size_w = tf.cast(input_size_w, tf.float32)
input_size_h = tf.cast(input_size_h, tf.float32)
bbox_loss_scale = 2.0 - 1.0 * label_xywh[:, :, :, :, 2:3] * label_xywh[:, :, :, :, 3:4] / (
input_size_w * input_size_h)
giou_loss = respond_bbox * bbox_loss_scale * (1 - giou) # giou_loss = (2 - bbox_area/image_area) * giou
Focal Loss
网格中的 anchor 是否包含目标,这是个逻辑回归问题。作者这里引入了 Focal Loss,给纯背景框的 loss 进行压缩,Focal loss 的作用可参考论文。
iou = self.bbox_iou(pred_xywh[:, :, :, :, np.newaxis, :], bboxes[:, np.newaxis, np.newaxis, np.newaxis, :, :])
# 找出与真实框 iou 值最大的预测框(这是为背景框有可能回归到目标来做准备的)
max_iou = tf.expand_dims(tf.reduce_max(iou, axis=-1), axis=-1)
# 如果最大的 iou 小于阈值,那么认为该预测框不包含物体,则为背景框
respond_bgd = (1.0 - respond_bbox) * tf.cast(max_iou < self.iou_loss_thresh, tf.float32)
# 计算 loss 权重,如果 anchor 不包含物体,且回归后仍然不包含物体,那个给这种背景框的 loss 更多的惩罚
conf_focal = self.focal(respond_bbox, pred_conf)
# 计算置信度的损失(我们希望假如该 anchor 包含物体,那么网络输出的预测框置信度为 1,无物体时则为 0。
conf_loss = conf_focal * (
respond_bbox * tf.nn.sigmoid_cross_entropy_with_logits(labels=respond_bbox, logits=conv_raw_conf)
+
respond_bgd * tf.nn.sigmoid_cross_entropy_with_logits(labels=respond_bbox, logits=conv_raw_conf)
)
分类损失
最后这个分类损失就没什么好说的了,采用的是二分类的交叉熵,即把所有类别的分类问题归结为是否属于这个类别,这样就把多分类看做是二分类问题。
prob_loss = respond_bbox * tf.nn.sigmoid_cross_entropy_with_logits(labels=label_prob, logits=conv_raw_prob)


